docstr: Test simple ResStage with different inputs.
code: def test_create_res_stage(self):
      """
      Test simple ResStage with different inputs.
      """
      for dim_in, dim_inner, dim_out in itertools.product(
          (4, 8, 16), (2, 4), (4, 8, 16)
      ):
          model = ResStage(
              res_blocks=nn.ModuleList(
                  [
                      ResBlock(
                          branch1_conv=nn.Conv3d(
                              dim_in, dim_out, kernel_size=(1, 1, 1)
                          )
                          if dim_in != dim_out
                          else None,
                          branch1_norm=nn.BatchNorm3d(
                              num_features=dim_out
                          )
                          if dim_in != dim_out
                          else None,
                          branch2=BottleneckBlock(
                              conv_a=nn.Conv3d(
                                  dim_in,
                                  dim_inner,
                                  kernel_size=[3, 1, 1],
                                  stride=[1, 1, 1],
                                  padding=[1, 0, 0],
                                  bias=False,
                              ),
                              norm_a=nn.BatchNorm3d(dim_inner),
                              act_a=nn.ReLU(),
                              conv_b=nn.Conv3d(
                                  dim_inner,
                                  dim_inner,
                                  kernel_size=[1, 3, 3],
                                  stride=[1, 1, 1],
                                  padding=[0, 1, 1],
                                  bias=False,
                              ),
                              norm_b=nn.BatchNorm3d(dim_inner),
                              act_b=nn.ReLU(),
                              conv_c=nn.Conv3d(
                                  dim_inner,
                                  dim_out,
                                  kernel_size=[1, 1, 1],
                                  stride=[1, 1, 1],
                                  padding=[0, 0, 0],
                                  bias=False,
                              ),
                              norm_c=nn.BatchNorm3d(dim_out),
                          ),
                          activation=nn.ReLU(),
                          branch_fusion=lambda x, y: x + y,
                      ),
                      ResBlock(
                          branch1_conv=None,
                          branch1_norm=None,
                          branch2=BottleneckBlock(
                              conv_a=nn.Conv3d(
                                  dim_out,
                                  dim_inner,
                                  kernel_size=[3, 1, 1],
                                  stride=[1, 1, 1],
                                  padding=[1, 0, 0],
                                  bias=False,
                              ),
                              norm_a=nn.BatchNorm3d(dim_inner),
                              act_a=nn.ReLU(),
                              conv_b=nn.Conv3d(
                                  dim_inner,
                                  dim_inner,
                                  kernel_size=[1, 3, 3],
                                  stride=[1, 1, 1],
                                  padding=[0, 1, 1],
                                  bias=False,
                              ),
                              norm_b=nn.BatchNorm3d(dim_inner),
                              act_b=nn.ReLU(),
                              conv_c=nn.Conv3d(
                                  dim_inner,
                                  dim_out,
                                  kernel_size=[1, 1, 1],
                                  stride=[1, 1, 1],
                                  padding=[0, 0, 0],
                                  bias=False,
                              ),
                              norm_c=nn.BatchNorm3d(dim_out),
                          ),
                          activation=nn.ReLU(),
                          branch_fusion=lambda x, y: x + y,
                      ),
                  ]
              )
          )
  
          # Test forwarding.
          for tensor in TestResStageTransform._get_inputs(dim_in):
              if tensor.shape[1] != dim_in:
                  with self.assertRaises(RuntimeError):
                      out = model(tensor)
                  continue
  
              out = model(tensor)
  
              input_shape = tensor.shape
              output_shape = out.shape
              output_shape_gt = (
                  input_shape[0],
                  dim_out,
                  input_shape[2],
                  input_shape[3],
                  input_shape[4],
              )
  
              self.assertEqual(
                  output_shape,
                  output_shape_gt,
                  "Output shape {} is different from expected shape {}".format(
                      output_shape, output_shape_gt
                  ),
              )
  
imports: import itertools
  import os
  import unittest
  import numpy
  import torch
  from pytorchvideo.models.head import ResNetBasicHead
  from pytorchvideo.models.net import Net
  from pytorchvideo.models.resnet import BottleneckBlock, create_acoustic_bottleneck_block, create_acoustic_resnet, create_bottleneck_block, create_res_block, create_res_stage, create_resnet, create_resnet_with_roi_head, ResBlock, ResStage, SeparableBottleneckBlock
  from pytorchvideo.models.stem import ResNetBasicStem
  from torch import nn
label_desc: Use image node_id 2 for calling the TestResStageTransform obj's underlying member method code with example usage: obj.test_create_res_stage() without return types