Is it possible to save a file from test_step() function?

You can aggregate test result in test_epoch_end:

def test_step(self, test_batch):
    x, y = test_batch
    logits = self.forward(x)
    loss = self.mean_squared_error_loss(logits.squeeze(-1), y.float())


    self.log('test_loss', loss)
    return {'test_loss': loss, "logits":logits, "labels": y}

def test_epoch_end(self, outputs):
    all_preds, all_labels = [], []
    for output in outputs:
        probs = list(output['logits'].cpu().detach().numpy()) # predicted values
        labels = list(output['labels'].flatten().cpu().detach().numpy())
        all_preds.extend(probs)
        all_labels.extend(labels)

    # you can calculate R2 here or save results as file
    r2 = ...

Note that this only works on a single GPU. If you are using multiple GPUs, you need some function to gather results from different GPUs.

To get model predictions, you need to add a predict_step() in the model class.

def predict_step(self, test_batch):
    x, y = test_batch
    logits = self.forward(x)
    return {'logits': logits, 'labels':y}

And run:

outputs = trainer.predict(model, test_loader, return_predictions=True)