Hugging Face's logo

Spaces:
wb-droid
/
MyMarioAI2
Sleeping

App Files Community
wb-droid commited on
Commit
b5ae145
·
1 Parent(s): 979ff77

minor fix

Browse files
Files changed (1) hide show
  1. MyMarioAI.py +1 -1
MyMarioAI.py CHANGED
@@ -34,7 +34,7 @@ def load_mario(env, save_dir):
34
 
35
  # Load from default chkpt.
36
  if Path(mario.default_chkpoint).is_file():
37
- dic = torch.load(mario.default_chkpoint)
38
  mario.net.load_state_dict(dic["model"])
39
  mario.exploration_rate = dic["exploration_rate"]
40
 
 
34
 
35
  # Load from default chkpt.
36
  if Path(mario.default_chkpoint).is_file():
37
+ dic = torch.load(mario.default_chkpoint, map_location=torch.device('cpu'))
38
  mario.net.load_state_dict(dic["model"])
39
  mario.exploration_rate = dic["exploration_rate"]
40