code: def test_crestereo(model_fn, model_mode, dev):
      set_rng_seed(0)
  
      model = model_fn().eval().to(dev)
  
      if model_mode == "scripted":
          model = torch.jit.script(model)
  
      img1 = torch.rand(1, 3, 64, 64).to(dev)
      img2 = torch.rand(1, 3, 64, 64).to(dev)
      iterations = 3
  
      preds = model(img1, img2, flow_init=None, num_iters=iterations)
      disparity_pred = preds[-1]
  
      # all the pyramid levels except the highest res make only half the number of iterations
      expected_iterations = (iterations // 2) * (
          len(model.resolutions) - 1
      )
      expected_iterations += iterations
      assert (
          len(preds) == expected_iterations
      ), "Number of predictions should be the number of iterations multiplied by the number of pyramid levels"
  
      assert disparity_pred.shape == torch.Size(
          [1, 2, 64, 64]
      ), f"Predicted disparity should have the same spatial shape as the input. Inputs shape {img1.shape[2:]}, Prediction shape {disparity_pred.shape[2:]}"
  
      assert all(
          d.shape == torch.Size([1, 2, 64, 64]) for d in preds
      ), "All predicted disparities are expected to have the same shape"
  
      # test a backward pass with a dummy loss as well
      preds = torch.stack(preds, dim=0)
      targets = torch.ones_like(preds, requires_grad=False)
      loss = torch.nn.functional.mse_loss(preds, targets)
  
      try:
          loss.backward()
      except Exception as e:
          assert (
              False
          ), f"Backward pass failed with an unexpected exception: {e.__class__.__name__} {e}"
  
      TM._assert_expected(
          disparity_pred, name=model_fn.__name__, atol=1e-2, rtol=1e-2
      )
  
imports: import pytest
  import test_models
  import torch
  from common_utils import cpu_and_cuda, set_rng_seed
  from torchvision.prototype import models
label_desc: Use image node_id 2 for calling a global function with example usage: test_crestereo(model_fn, model_mode, dev) without return types