fc94cf593f218f4c3e715d8a2d58b324c2828007,VAEs/vq-vae/trainers/train_prior.py,TrainerPrior,forward_pass,#TrainerPrior#Any#Any#,83

Before Change

	def forward_pass(self, img_var, labels):
		enc_indices, quantized = self.base_model(img_var, return_encoding_indices=True, test=True)
		labels_var = nn.Variable(labels.shape)
		if isinstance(labels, nn.NdArray):
			labels_var.data  = labels
		else:
			labels_var.d = labels
		labels_var = F.one_hot(labels_var, shape=(self.num_classes,))
		enc_recon = self.prior(quantized, labels_var)
		loss = F.mean(F.softmax_cross_entropy(enc_recon, enc_indices))

After Change

	def forward_pass(self, img_var, labels):
		enc_indices, quantized = self.base_model(img_var, return_encoding_indices=True)

		if self.dataset_name == "imagenet":
			labels = nn.Variable(labels.shape).apply(data=labels)
		else:
			labels = nn.Variable.from_numpy_array(labels)
		labels = F.one_hot(labels, shape=(self.num_classes,))
		enc_recon = self.prior(quantized, labels)
		loss = F.mean(F.softmax_cross_entropy(enc_recon, enc_indices))

Instances