def _compute_unrolled_model(self, input_source,input_target, domain_label_source,domain_label_target, eta, network_optimizer):
loss_source = self.model._loss(input_source, domain_label_source)

loss_target = self.model._loss(input_target,domain_label_target)

loss=loss_source+loss_target

theta = _concat(self.model.parameters()).data

try:

  moment = _concat(network_optimizer.state[v]['momentum_buffer'] for v in self.model.parameters()).mul_(self.network_momentum)

except:

  moment = torch.zeros_like(theta)

dtheta = _concat(torch.autograd.grad(loss, self.model.parameters(),retain_graph=True)).data + self.network_weight_decay*theta

unrolled_model = self._construct_model_from_theta(theta.sub(eta, moment+dtheta))

return unrolled_model

def _construct_model_from_theta(self, theta):
model_new = self.model.new()

model_dict = self.model.state_dict()

params, offset = {}, 0

for k, v in self.model.named_parameters():

  v_length = np.prod(v.size())

  params[k] = theta[offset: offset+v_length].view(v.size())

  offset += v_length

assert offset == len(theta)

model_dict.update(params)

model_new.load_state_dict(model_dict)

return model_new.cuda()