flatten predictions from batches

[pamparana34]

I am currently using the trainer predict function for doing some inference and this might be a bit silly but I am calling it as:

predictions = trainer.predict(model=model, datamodule=datamodule, ckpt_path=config.ckpt_path)

Now, the issue is that the dataloader has some greater than 1 batch size set (typically 256), so that I can use the GPU effectively to run inference. So the returned object is a list of list of pytorch tensors and this makes it difficult to output this in desired format.

So, my predict function is something like:

def predict_step(self, batch, batch_idx, dataloader_idx=None):
        x, y, _ = batch
        return self(x)

I tried hooking onto on_predict_epoch_endas:

def on_predict_epoch_end(self, results: List[Any]):
        flattened_tensors = list()

        for l in results:
            flattened_tensors += torch.vstack(l).tolist()
            
        return flattened_tensors

However, this does not seem to be able to influence the returned tensor. I think nothing can be returned from this method but not sure why it accepts the results param.

Usually there is something like train_epoch_end where we can return values and I wonder how I can do the equivalent for the prediction steps.

For various design reasons, I cannot do this on the side of the caller. Is there a way to do this from within the model code through a hook?