Update train.py

train.py

@@ -31,7 +31,7 @@ if __name__ == "__main__":
     pretrained_dict = torch.load("model_data/ssd_weights.pth")
     pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
     model_dict.update(pretrained_dict)
-    model.load_state_dict(pretrained_dict)
+    model.load_state_dict(model_dict)
     print('Finished!')
 
     net = model