From 1e791f3afc502873ac0eb413cdff4140a7c46bdd Mon Sep 17 00:00:00 2001 From: Rasmus Oersoe Date: Tue, 5 Sep 2023 12:38:51 +0200 Subject: [PATCH] Add check for torch threads in deployer --- src/graphnet/deployment/i3modules/deployer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/graphnet/deployment/i3modules/deployer.py b/src/graphnet/deployment/i3modules/deployer.py index 5db7176fc..2228005f8 100644 --- a/src/graphnet/deployment/i3modules/deployer.py +++ b/src/graphnet/deployment/i3modules/deployer.py @@ -58,8 +58,11 @@ def __init__( """ # This makes sure that one worker cannot access more # than 1 core's worth of compute. - torch.set_num_threads(1) - torch.set_num_interop_threads(1) + + if torch.get_num_interop_threads() > 1: + torch.set_num_interop_threads(1) + if torch.get_num_threads() > 1: + torch.set_num_threads(1) # Check if isinstance(graphnet_modules, list): self._modules = graphnet_modules