diff --git a/setup.py b/setup.py
index b2a45ff..7a848be 100644
--- a/setup.py
+++ b/setup.py
@@ -59,8 +59,8 @@ def get_string(string, rel_path="src/stepcount/__init__.py"):
         "scikit-learn==1.1.1",
         "imbalanced-learn==0.9.1",
         "hmmlearn==0.3.*",
-        "torch==1.13.*",
-        "torchvision==0.14.*",
+        "torch>=1.13,<3",
+        "torchvision>=0.14,<1",
         "transforms3d==0.4.*",
         "numba==0.58.*",
         "matplotlib==3.7.*",
diff --git a/src/stepcount/models.py b/src/stepcount/models.py
index e8f753b..73e1542 100644
--- a/src/stepcount/models.py
+++ b/src/stepcount/models.py
@@ -441,7 +441,7 @@ def fit(self, X, Y, groups=None):
         model.to(self.device)
 
         sslmodel.train(model, train_loader, val_loader, self.device, class_weights, weights_path=self.weights_path)
-        model.load_state_dict(torch.load(self.weights_path, self.device))
+        model.load_state_dict(torch.load(self.weights_path, map_location=self.device))
 
         if self.verbose:
             print('Training HMM')