class_docstr: Test NonLocal layers.
docstr: Test non-local block2d.
code: def test_non_local_block2d(self):
      """
      Test non-local block2d.
      """
      in_channels = 2
      inter_channels = None
      sub_sample = True
      bn_layer = True
      net = NONLocalBlock2D(
          in_channels=in_channels,
          inter_channels=inter_channels,
          sub_sample=sub_sample,
          bn_layer=bn_layer,
      )
      input_tensor = torch.randn(1, 2, 4, 4)
      output_tensor = net(input_tensor)
      self.assertTrue(output_tensor.shape == torch.Size([1, 2, 4, 4]))
  
file_path: towhee/tests/unittests/models/layers/test_non_local.py
imports: import unittest
  import torch
  from towhee.models.layers.non_local import _NonLocalBlockND, NONLocalBlock1D, NONLocalBlock2D, NONLocalBlock3D
label_desc: Use image node_id 3 for calling the TestNonLocal obj's underlying member method code with example usage: obj.test_non_local_block2d() without return types