question: Saving one single inference example per validation stage
context: 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
answer: 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)
  ...