Pytorch implementation of meta updating about maml-pytorch HOT 6 CLOSED

		weights_origin = copy.deepcopy(self.net.parameters())
		for i in range(batchsz): # batchsz==self.meta_batchsz
			pred_q = self.net(query_x[i], bns=None, training=training)
			pred_q = F.softmax(pred_q, dim=1).argmax(dim=1)
			correct = torch.eq(pred_q, query_y[i]).sum().item()
			corrects[0] = corrects[0] + correct

			# 1. run the i-th task and compute loss for k=0
			self.inner_optim.zero_grad()
			pred = self.net(support_x[i])
			loss = F.cross_entropy(pred, support_y[i])
			loss.backward()
			nn.utils.clip_grad_norm_(self.net.parameters(), 5)
			self.inner_optim.step()

			pred_q = self.net(query_x[i], bns=None, training=training)
			loss_q = F.cross_entropy(pred_q, query_y[i])

			pred_q = F.softmax(pred_q, dim=1).argmax(dim=1)
			correct = torch.eq(pred_q, query_y[i]).sum().item()
			corrects[1] = corrects[1] + correct

			for k in range(1, self.K):
				# 1. run the i-th task and compute loss for k=1~K-1
				self.inner_optim.zero_grad()
				pred = self.net(support_x[i], bns=None, training=training)
				loss = F.cross_entropy(pred, support_y[i])
				loss.backward()
				nn.utils.clip_grad_norm_(self.net.parameters(), 5)
				self.inner_optim.step()

				pred_q = self.net(query_x[i], bns=None, training=training)
				loss_q = F.cross_entropy(pred_q, query_y[i])
				pred_q = F.softmax(pred_q, dim=1).argmax(dim=1)
				correct = torch.eq(pred_q, query_y[i]).sum().item() # convert to numpy
				corrects[k+1] = corrects[k+1] + correct

			# 4. record last step's loss for task i
			losses_q.append(loss_q)
			for p, p_o in zip(self.net.parameters(), weights_origin):
				p.data.copy_(p_o)

		# end of all tasks
		# sum over all losses across all tasks
		loss_q = torch.stack(losses_q).sum(0)