File size: 282 Bytes
f348660
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
from safetensors.torch import load_file


class BaseModel(torch.nn.Module):
    def load(self, path):
        """Load model from file.

        Args:
            path (str): file path
        """
        self.load_state_dict(load_file(path, device="cpu"), strict=True)