Spaces:

aki-0421
/

character-360

Build error

character-360 / annotator /midas /base_model.py

aki-0421

F: add

a3a3ae4 unverified almost 2 years ago

391 Bytes

# -*- coding: utf-8 -*-
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)

	# -- coding: utf-8 --
	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)