    def test_create_stem_with_callable(self):
        """
        Test builder `create_res_basic_stem` with callable inputs.
        """
        for (pool, activation, norm) in itertools.product(
            (nn.AvgPool3d, nn.MaxPool3d, None),
            (nn.ReLU, nn.Softmax, nn.Sigmoid, None),
            (nn.BatchNorm3d, None),
        ):
            model = create_resnet_basic_3d_module(
                in_channels=3,
                out_channels=64,
                pool=pool,
                activation=activation,
                norm=norm,
            )
            model_gt = ResNetBasic3DModule(
                conv=nn.Conv3d(
                    3,
                    64,
                    kernel_size=[3, 7, 7],
                    stride=[1, 2, 2],
                    padding=[1, 3, 3],
                    bias=False,
                ),
                norm=None if norm is None else norm(64),
                activation=None if activation is None else activation(),
                pool=None
                if pool is None
                else pool(kernel_size=[1, 3, 3], stride=[1, 2, 2], padding=[0, 1, 1]),
            )

            model.load_state_dict(
                model_gt.state_dict(), strict=True
            )  # explicitly use strict mode.

            # Test forwarding.
            for input_tensor in TestResNetBasic3DModule._get_inputs():
                with torch.no_grad():
                    if input_tensor.shape[1] != 3:
                        with self.assertRaises(RuntimeError):
                            output_tensor = model(input_tensor)
                        continue
                    else:
                        output_tensor = model(input_tensor)
                        output_tensor_gt = model_gt(input_tensor)
                self.assertEqual(
                    output_tensor.shape,
                    output_tensor_gt.shape,
                    f"Output shape {output_tensor.shape} is different from expected shape {output_tensor_gt.shape}"
                )
                self.assertTrue(
                    np.allclose(output_tensor.numpy(), output_tensor_gt.numpy())
                )
Test builder `create_res_basic_stem` with callable inputs.
tests/unittests/models/layers/test_resnet_basic_3d_module.py