def load_model(model_type, checkpoint_file):
checkpoint = torch.load(checkpoint_file)
args = checkpoint['args']
if model_type == 'teacher':
model = DCTeacherNet(args).to(args.device)
else:
model = DCStudentNet(args).to(args.device)
model.load_state_dict(checkpoint['model'])
return model
def build_model(args):
if args.model_type == 'teacher':
model = DCTeacherNet(args).to(args.device)
if args.model_type == 'student':
model = DCStudentNet(args).to(args.device)
return model
def build_model(args):
modelT = DCTeacherNet(args).to(args.device)
modelS = DCStudentNet(args).to(args.device)
return modelT,modelS