Hugging Face's logo

Spaces:
danielcd99
/
MNIST-classifier
Runtime error

App Files Community
1
danielcd99 commited on
Commit
e7f55d6
·
1 Parent(s): b025479

Changed for cpu

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -39,7 +39,7 @@ class LeNet(nn.Module):
39
 
40
  # LOADING MODEL
41
  model = LeNet().to(device)
42
- model.load_state_dict(torch.load("model_mnist.pth", map_location=torch.device('cuda')))
43
 
44
 
45
  def predict(input):
 
39
 
40
  # LOADING MODEL
41
  model = LeNet().to(device)
42
+ model.load_state_dict(torch.load("model_mnist.pth", map_location=torch.device('cpu')))
43
 
44
 
45
  def predict(input):