code: def run_test(batch_size=64, train_steps=200, device="cuda:0"):
      class _TestBatch(Dataset):
          def __init__(self, transforms):
              self.transforms = transforms
  
          def __getitem__(self, _unused_id):
              im, seg = create_test_image_2d(
                  128, 128, noise_max=1, num_objs=4, num_seg_classes=1
              )
              seed = np.random.randint(2147483647)
              self.transforms.set_random_state(seed=seed)
              im = self.transforms(im)
              self.transforms.set_random_state(seed=seed)
              seg = self.transforms(seg)
              return im, seg
  
          def __len__(self):
              return train_steps
  
      net = UNet(
          spatial_dims=2,
          in_channels=1,
          out_channels=1,
          channels=(4, 8, 16, 32),
          strides=(2, 2, 2),
          num_res_units=2,
      ).to(device)
  
      loss = DiceLoss(sigmoid=True)
      opt = torch.optim.Adam(net.parameters(), 1e-2)
      train_transforms = Compose(
          [
              EnsureChannelFirst(channel_dim="no_channel"),
              ScaleIntensity(),
              RandSpatialCrop((96, 96), random_size=False),
              RandRotate90(),
          ]
      )
  
      src = DataLoader(
          _TestBatch(train_transforms),
          batch_size=batch_size,
          shuffle=True,
      )
  
      net.train()
      epoch_loss = 0
      step = 0
      for img, seg in src:
          step += 1
          opt.zero_grad()
          output = net(img.to(device))
          step_loss = loss(output, seg.to(device))
          step_loss.backward()
          opt.step()
          epoch_loss += step_loss.item()
      epoch_loss /= step
  
      return epoch_loss, step
  
file_path: MONAI/tests/test_integration_determinism.py
imports: from __future__ import annotations
  import unittest
  import numpy
  import torch
  from torch.utils.data import DataLoader, Dataset
  from monai.data import create_test_image_2d
  from monai.losses import DiceLoss
  from monai.networks.nets import UNet
  from monai.transforms import Compose, EnsureChannelFirst, RandRotate90, RandSpatialCrop, ScaleIntensity
  from monai.utils import set_determinism
  from tests.utils import DistTestCase, TimedCall
label_desc: Use image node_id 1 for calling a global function with example usage: run_test(batch_size, train_steps, device) and returns: epoch_loss,  step, im,  seg, train_steps