Method predict_step

src/flair/task_module.py:206–213

(self, batch, batch_idx, dataloader_idx=0)

        return preds, targets

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        if self.use_metadata == True:
            logits = self.forward(batch["img"], batch["mtd"])
        else:
            logits = self.forward(batch["img"], '')
        proba = torch.softmax(logits, dim=1)
        batch["preds"] =  torch.argmax(proba, dim=1)
        return batch