docstr: Test builder `create_res_stage` with callable inputs.
code: def test_create_res_stage_with_callable(self):
      """
      Test builder `create_res_stage` with callable inputs.
      """
      dim_in, dim_inner, dim_out = 32, 16, 64
      for norm, activation in itertools.product(
          (nn.BatchNorm3d, None), (nn.ReLU, nn.Sigmoid, None)
      ):
          model = create_res_stage(
              depth=2,
              dim_in=dim_in,
              dim_inner=dim_inner,
              dim_out=dim_out,
              bottleneck=create_bottleneck_block,
              conv_a_kernel_size=(3, 1, 1),
              conv_a_stride=(1, 1, 1),
              conv_a_padding=(1, 0, 0),
              conv_b_kernel_size=(1, 3, 3),
              conv_b_stride=(1, 1, 1),
              conv_b_padding=(0, 1, 1),
              conv_b_num_groups=1,
              conv_b_dilation=(1, 1, 1),
              norm=norm,
              norm_eps=1e-5,
              norm_momentum=0.1,
              activation=activation,
          )
          model_gt = ResStage(
              res_blocks=nn.ModuleList(
                  [
                      ResBlock(
                          branch1_conv=nn.Conv3d(
                              dim_in,
                              dim_out,
                              kernel_size=(1, 1, 1),
                              bias=False,
                          )
                          if dim_in != dim_out
                          else None,
                          branch1_norm=None
                          if norm is None
                          else norm(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=None
                              if norm is None
                              else norm(dim_inner),
                              act_a=None
                              if activation is None
                              else activation(),
                              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=None
                              if norm is None
                              else norm(dim_inner),
                              act_b=None
                              if activation is None
                              else activation(),
                              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=None
                              if norm is None
                              else norm(dim_out),
                          ),
                          activation=None
                          if activation is None
                          else activation(),
                          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=None
                              if norm is None
                              else norm(dim_inner),
                              act_a=None
                              if activation is None
                              else activation(),
                              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=None
                              if norm is None
                              else norm(dim_inner),
                              act_b=None
                              if activation is None
                              else activation(),
                              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=None
                              if norm is None
                              else norm(dim_out),
                          ),
                          activation=None
                          if activation is None
                          else activation(),
                          branch_fusion=lambda x, y: x + y,
                      ),
                  ]
              )
          )
          model.load_state_dict(
              model_gt.state_dict(), strict=True
          )  # explicitly use strict mode.
  
          # Test forwarding.
          for tensor in TestResStageTransform._get_inputs(
              dim_in=dim_in
          ):
              with torch.no_grad():
                  if tensor.shape[1] != 32:
                      with self.assertRaises(RuntimeError):
                          out = model(tensor)
                      continue
  
                  out = model(tensor)
                  out_gt = model_gt(tensor)
  
              self.assertEqual(
                  out.shape,
                  out_gt.shape,
                  "Output shape {} is different from expected shape {}".format(
                      out.shape, out_gt.shape
                  ),
              )
              self.assertTrue(np.allclose(out.numpy(), out_gt.numpy()))
  
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 3 for calling the TestResStageTransform obj's underlying member method code with example usage: obj.test_create_res_stage_with_callable() without return types