load with snapshot_download

app.py

@@ -3,6 +3,7 @@ import numpy as np
 import torch
 import torchvision
 from torch import nn
+from huggingface_hub import snapshot_download
 
 class LeNet(nn.Module):
     def __init__(self):
@@ -35,8 +36,9 @@ class LeNet(nn.Module):
         return nn.functional.softmax(out[0], dim = 0)
 
 lenet = LeNet()
-lenet.load_state_dict(torch.load('stanimirovb/ibob-lenet-v1/lenet-v1.pth', map_location='cpu'))
 
+lenet_pt = snapshot_download('stanimirovb/ibob-lenet-v1') + '/lenet-v1.pth'
+lenet.load_state_dict(torch.load(lenet_pt, map_location='cpu'))
 
 resize = torchvision.transforms.Resize((28, 28), antialias=True)
 def on_submit(img):