jerilseb commited on
Commit
9a4023b
1 Parent(s): 54f8640

update stuff

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -31,11 +31,17 @@ state_dict = torch.load('model.pth', map_location='cpu')
31
  model.load_state_dict(state_dict, strict=False)
32
  model.eval()
33
 
 
 
 
 
 
 
34
  def predict(image):
35
  image = image['composite']
36
- image = transforms.Resize((28, 28))(image)
37
  print(image)
38
-
39
  x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
40
 
41
  with torch.no_grad():
 
31
  model.load_state_dict(state_dict, strict=False)
32
  model.eval()
33
 
34
+ transform = transforms.Compose([
35
+ transforms.Resize((28, 28)),
36
+ transforms.ToTensor(),
37
+ transforms.Normalize((0.5,), (0.5,))
38
+ ])
39
+
40
  def predict(image):
41
  image = image['composite']
42
+ image = transform(image)
43
  print(image)
44
+
45
  x = torch.tensor(im, dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.
46
 
47
  with torch.no_grad():