Load the checkpoint into the cpu

model.py

@@ -25,7 +25,7 @@ def get_model(checkpoint_path):
     return model
 
 def get_checkpoint(checkpoint_path):
-    ckpt = torch.load('checkpoint.ckpt')
+    ckpt = torch.load('checkpoint.ckpt', map_location=torch.device('cpu'))
 
     fixed_state_dict = collections.OrderedDict()