Saving one single inference example per validation stage

[leoauri]

I would like to save one single inference example per validation stage. To this end I came up with:

class Model(pl.LightningModule):

    ...

    def validation_epoch_end(self, validation_step_outputs):
        # self.trainer.log_dir is not set during fast_dev_test
        if self.trainer.log_dir is not None:
            x, y = (d.unsqueeze(0) for d in self.trainer.datamodule.valid_set.dataset_pair())
            y_hat = self.model(x)
            x, y, y_hat = (d.squeeze().cpu().numpy() for d in (x, y, y_hat))

            save_dir = Path(self.trainer.log_dir) / 'wavs'
            save_dir.mkdir(exist_ok=True)

            soundfile.write(save_dir/f'{self.global_step:06}_input.wav',
                    x, 44100, subtype='PCM_24')
            soundfile.write(save_dir/f'{self.global_step:06}_output.wav',
                    y_hat, 44100, subtype='PCM_24')
            soundfile.write(save_dir/f'{self.global_step:06}_target.wav',
                    y, 44100, subtype='PCM_24')

This works OK when running on a single device (CPU) but when running on GPUs on a slurm cluster I get:

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

I guess because the data has not been sent to the same device as the model.

This seems like a problem that must have been solved before. Can anyone suggest a typical pytorch-lightning way of doing this?

Thanks for your input
Leo

[rohitgr7]

here:

x, y = (d.unsqueeze(0) for d in self.trainer.datamodule.valid_set.dataset_pair())

I think this is something not part of the dataloader, so PL won't move to the device automatically.
you can do:

x = x.to(self.device)
y_hat = self.model(x)
...

[leoauri]

Great, thank you!