Update app.py

app.py

@@ -48,6 +48,7 @@ DesignModernityModel = torch.load("DesignModernityModel.pt")
 DesignModernityModel.eval() # set state of the model to inference
 
 LABELS = ['2000-2004', '2006-2008', '2009-2011', '2012-2015', '2016-2018']
+n_labels = len(LABELS)
 
 MEAN = [0.485, 0.456, 0.406]
 STD = [0.229, 0.224, 0.225]
@@ -61,7 +62,7 @@ def classifyCar(im):
   im = carTransforms(im).unsqueeze(0)  # transform and add batch dimension
   with torch.no_grad():
     scores = torch.nn.functional.softmax(DesignModernityModel(im)[0])
-  return {LABELS[i]: float(scores[i]) for i in range(2)}
+  return {LABELS[i]: float(scores[i]) for i in range(n_labels)}
 
 #examples = [[example_img.jpg], [example_img2.jpg]]  # must be uploaded in repo