def test_step(self, batch, batch_idx):
vilt_utils.set_task(self)
output = self(batch)
ret = dict()
if self.hparams.config["loss_names"]["vqa"] > 0:
ret.update(objectives.vqa_test_step(self, batch, output))
return ret
def validation_step(self, batch, batch_idx):
vilt_utils.set_task(self)
output = self(batch)
def training_step(self, batch, batch_idx):
vilt_utils.set_task(self)
output = self(batch)
total_loss = sum([v for k, v in output.items() if "loss" in k])
return total_loss