Merge pull request #11 from RMeli/jit

Compile models to TorchScript
RMeli · Nov 17, 2021 · 563fe78 · 563fe78
2 parents 2e4457f + 4593d45
commit 563fe78
Show file tree

Hide file tree

Showing 4 changed files with 352 additions and 96 deletions.
diff --git a/gnina/inference.py b/gnina/inference.py
@@ -151,9 +151,7 @@ def inference(args):
     affinity: bool = True if args.affinity_pos is not None else False
 
     # Create model
-    model = models.models_dict[args.model](test_loader.dims, affinity=affinity).to(
-        device
-    )
+    model = models.models_dict[(args.model, affinity)](test_loader.dims).to(device)
 
     # Load checkpoint
     checkpoint = torch.load(args.checkpoint, map_location=device)