Bugfix

src/models/train_model.py

@@ -45,8 +45,8 @@ def get_model(name: str, triples, inverse_triples, regularizer):
         "ConvE": 0.1,
         "DistMult": 0.0,
     }
+    model = model_dict[name]
     if name == "TransE":
-        model = model_dict[name]
         model.loss.margin = 2.0
     return model, lr_dict[name], epochs_dict[name], bs_dict[name], ls_dict[name]