Spaces:
Sleeping
Sleeping
load with snapshot_download
Browse files
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):
|