code: def test_index(self, mock_to_file):
      callback = imaging.ImagingCallback().to_file("test", index=10)
      image = torch.zeros(11, 1, 1, 1)
      callback.on_batch = lambda _: image
      callback.process("state")
      mock_to_file().assert_called_once_with(ANY, 10, "state")
      self.assertTrue(
          (mock_to_file().call_args_list[0][0][0] == image[10]).all()
      )
      mock_to_file().reset_mock()
  
      callback = imaging.ImagingCallback().to_file(
          "test", index=[2, 5, 10]
      )
      image = torch.rand(11, 1, 1, 1)
      callback.on_batch = lambda _: image
      callback.process("state")
      self.assertTrue(mock_to_file().call_count == 3)
      self.assertTrue(
          (mock_to_file().call_args_list[0][0][0] == image[2]).all()
      )
      self.assertTrue(
          (mock_to_file().call_args_list[1][0][0] == image[5]).all()
      )
      self.assertTrue(
          (mock_to_file().call_args_list[2][0][0] == image[10]).all()
      )
  
imports: from unittest import TestCase
  from mock import MagicMock, patch, ANY
  import torch
  import torchbearer
  import torchbearer.callbacks.imaging
  import matplotlib.pyplot
label_desc: Use image node_id 6 for calling the TestImagingCallback obj's underlying member method code with example usage: obj.test_index(mock_to_file) without return types