docstr: Test builder `create_nonlocal`.
docstr_tok: ["Test","builder","`","create_nonlocal","`","."]
code: def test_nonlocal_builder(self):
      """
      Test builder `create_nonlocal`.
      """
      for (
          dim_in,
          dim_inner,
          pool_size,
          norm,
          instantiation,
      ) in itertools.product(
          (4, 8),
          (2, 4),
          ((1, 1, 1), (2, 2, 2)),
          (None, nn.BatchNorm3d),
          ("dot_product", "softmax"),
      ):
          conv_theta = nn.Conv3d(
              dim_in, dim_inner, kernel_size=1, stride=1, padding=0
          )
          conv_phi = nn.Conv3d(
              dim_in, dim_inner, kernel_size=1, stride=1, padding=0
          )
          conv_g = nn.Conv3d(
              dim_in, dim_inner, kernel_size=1, stride=1, padding=0
          )
          conv_out = nn.Conv3d(
              dim_inner, dim_in, kernel_size=1, stride=1, padding=0
          )
          if norm is None:
              norm_model = None
          else:
              norm_model = norm(num_features=dim_in)
          if isinstance(pool_size, Iterable) and any(
              size > 1 for size in pool_size
          ):
              pool_model = nn.MaxPool3d(
                  kernel_size=pool_size,
                  stride=pool_size,
                  padding=[0, 0, 0],
              )
          else:
              pool_model = None
  
          model = create_nonlocal(
              dim_in=dim_in,
              dim_inner=dim_inner,
              pool_size=pool_size,
              instantiation=instantiation,
              norm=norm,
          )
  
          model_gt = NonLocal(
              conv_theta=conv_theta,
              conv_phi=conv_phi,
              conv_g=conv_g,
              conv_out=conv_out,
              pool=pool_model,
              norm=norm_model,
              instantiation=instantiation,
          )
          model.load_state_dict(
              model_gt.state_dict(), strict=True
          )  # explicitly use strict mode.
  
          # Test forwarding.
          for input_tensor in TestNonlocal._get_inputs(
              input_dim=dim_in
          ):
              with torch.no_grad():
                  if input_tensor.shape[1] != dim_in:
                      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,
                  "Output shape {} is different from expected shape {}".format(
                      output_tensor.shape, output_tensor_gt.shape
                  ),
              )
              self.assertTrue(
                  np.allclose(
                      output_tensor.numpy(), output_tensor_gt.numpy()
                  )
              )
  
file_path: pytorchvideo/tests/test_layers_nonlocal_net.py
imports: import itertools
  import unittest
  from typing import Iterable
  import numpy
  import torch
  from pytorchvideo.layers.nonlocal_net import create_nonlocal, NonLocal
  from torch import nn
label_desc: Use image node_id 3 for calling the TestNonlocal obj's underlying member method code with example usage: obj.test_nonlocal_builder() without return types