fix squeezenet load

app.py

@@ -138,14 +138,15 @@ def infer(audio_path: str, log_name: str):
                 weight_path=f"{MODEL_DIR}/{log_name}.pt",
             )
 
+        input_size = eval_net.get_input_size()
+        embeded_input = embed(input, input_size)
+        output = list(eval_net.forward(embeded_input))
+
     except Exception as e:
         return f"{e}", None
 
-    input_size = eval_net.get_input_size()
-    embeded_input = embed(input, input_size)
-    output = list(eval_net.forward(embeded_input))
-    outputs = []
     index = 0
+    outputs = []
     for y in output:
         preds = list(y.T)
         for pred in preds:

model.py

@@ -56,9 +56,15 @@ class EvalNet:
             torch.load(weight_path)
             if torch.cuda.is_available()
             else torch.load(weight_path, map_location="cpu")
-        )  # self.model.load_state_dict(checkpoint, False)
-        self.model.load_state_dict(checkpoint["model"], False)
-        self.classifier.load_state_dict(checkpoint["classifier"], False)
+        )
+
+        if self.type == "squeezenet":
+            self.model.load_state_dict(checkpoint, False)
+
+        else:
+            self.model.load_state_dict(checkpoint["model"], False)
+            self.classifier.load_state_dict(checkpoint["classifier"], False)
+
         if torch.cuda.is_available():
             self.model = self.model.cuda()
             self.classifier = self.classifier.cuda()