House-keeping

src/models/__init__.py

@@ -48,7 +48,8 @@ def load_model(path: Path, device: torch.device) -> Tuple[nn.Module, Dict[str, A
     model = create_model(
         type=hyperparameters["model_type"],
         bin_count=hyperparameters["bin_count"],
-    ).to(device)
+        device=device,
+    )
     model.load_state_dict(torch.load(model_path))
     model.eval()
     logging.info(f"Parameter count: {sum(p.numel() for p in model.parameters())}")