stanimirovb commited on
Commit
1364f9f
·
verified ·
1 Parent(s): fae33e7

load with snapshot_download

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -3,6 +3,7 @@ import numpy as np
3
  import torch
4
  import torchvision
5
  from torch import nn
 
6
 
7
  class LeNet(nn.Module):
8
  def __init__(self):
@@ -35,8 +36,9 @@ class LeNet(nn.Module):
35
  return nn.functional.softmax(out[0], dim = 0)
36
 
37
  lenet = LeNet()
38
- lenet.load_state_dict(torch.load('stanimirovb/ibob-lenet-v1/lenet-v1.pth', map_location='cpu'))
39
 
 
 
40
 
41
  resize = torchvision.transforms.Resize((28, 28), antialias=True)
42
  def on_submit(img):
 
3
  import torch
4
  import torchvision
5
  from torch import nn
6
+ from huggingface_hub import snapshot_download
7
 
8
  class LeNet(nn.Module):
9
  def __init__(self):
 
36
  return nn.functional.softmax(out[0], dim = 0)
37
 
38
  lenet = LeNet()
 
39
 
40
+ lenet_pt = snapshot_download('stanimirovb/ibob-lenet-v1') + '/lenet-v1.pth'
41
+ lenet.load_state_dict(torch.load(lenet_pt, map_location='cpu'))
42
 
43
  resize = torchvision.transforms.Resize((28, 28), antialias=True)
44
  def on_submit(img):