Bugfix

Models.py

@@ -134,13 +134,13 @@ class VisionModel(nn.Module):
 			from safetensors.torch import load_file
 			resume = load_file(Path(path) / 'model.safetensors', device='cpu')
 		else:
-			resume = torch.load(Path(path) / 'model.pt', map_location=torch.device('cpu'))
+			resume = torch.load(Path(path) / 'model.pt', map_location=torch.device('cpu'))['model']
 
 		model_classes = VisionModel.__subclasses__()
 		model_cls = next(cls for cls in model_classes if cls.__name__ == config['class'])
 
 		model = model_cls(**{k: v for k, v in config.items() if k != 'class'})
-		model.load(resume['model'])
+		model.load(resume)
 		if device is not None:
 			model = model.to(device)