Alternate "prediction" loops

[kaboroevich]

I am hoping to implement an additional loop through the Trainer in order to leverage Lightnings automagic handling of dataloaders and GPUs. Specifically, I want to run batches through the attribution methods from Captum. My first attempt was to hijack the predition_step of the LightningModule with a simple class attribute self.calculate_attributes switch:

    def predict_step(self, batch, batch_idx):
        if self.calculate_attributes:
            return self.attribution_step(batch, batch_idx)
        else:
            data, target = batch
            return self.model(data)
    
    def attribution_step(self, batch, batch_idx):
        data, target = batch
        batch_size = data.shape[0]
        baselines = torch.zeros_like(data)
        attribution = self.explainer.attribute(data, baselines, target=target, internal_batch_size=batch_size)
        return attribution, target

But this has run into issues because gradients are required, and, I believe, the prediction loop disables them. I tried to get around with the @torch.enable_grad() decorator and explicit .requires_grad=True calls, but that is not working. For every attempt I get the error RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn. I have no issue running the explainers on their own on the same data. I've been looking at the depreciated Loops interface but haven't made any progress there either.

Is there a proper way to implement this? Any suggestions would be appreciated.

[kaboroevich]

I have also tested calling attribution_step within training_step and it does not raise an error.

[kaboroevich]

Setting Trainer(inference_mode=False) was the answer

#15925 #15765