Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import os
 import torch
 import torch.nn.functional as F
-from mobilenet import MobileNetLarge3D
+from mobilenet import MobileNetLarge3D, MobileNetSmall3D
 from torchvision.io import read_video
 import time
 
@@ -38,8 +38,8 @@ def call_pitch(pitch):
     pitch_tensor = (read_video(pitch,pts_unit='sec')[0].permute(-1,0,1,2)).unsqueeze(0)/255
     pitch_tensor = (pitch_tensor-mean)/std #normalize the pitch tensor
     video_length = pitch_tensor.shape[2]/15
-    model = MobileNetLarge3D()
-    model.load_state_dict(torch.load('weights/MobileNetLarge.pth',map_location=torch.device('cpu')))
+    model = MobileNetSmall3D()
+    model.load_state_dict(torch.load('weights/MobileNetSmall.pth',map_location=torch.device('cpu')))
     model.eval()
 
     #run the model