Update model.py

model.py

@@ -21,4 +21,4 @@ class Model(torch.nn.Module):
         torch.save(self.state_dict(), path)
 
     def load(self, path):
-        self.load_state_dict(torch.load(path, weights_only=True))
+        self.load_state_dict(torch.load(path, weights_only=True, map_location=torch.device('cpu')))