Fix loading checkpoints to specified device
@@ -393,7 +393,11 @@ class KgeModel(KgeBase):
"""
+ if device:
+ checkpoint = torch.load(filename, map_location=device)
+ else:
checkpoint = torch.load(filename, map_location="cpu")
+
original_config = checkpoint["config"]
config = Config() # round trip to handle deprecated configs
config.load_options(original_config.options)
