dreamdrone / midas /base_model.py
imsuperkong's picture
Upload 31 files
dc47947
raw
history blame
367 Bytes
import torch
class BaseModel(torch.nn.Module):
def load(self, path):
"""Load model from file.
Args:
path (str): file path
"""
parameters = torch.load(path, map_location=torch.device('cpu'))
if "optimizer" in parameters:
parameters = parameters["model"]
self.load_state_dict(parameters)