code: def test_input_across_dims(
      self,
      N: int,
      sz: int,
      V: int,
      D: int,
      mode: str,
  ):
      emb = nn.EmbeddingBag(
          num_embeddings=V, embedding_dim=D, mode=mode
      )
  
      sizes = torch.randint(low=1, high=sz + 1, size=(N,))
      offsets = torch.LongTensor(
          [0] + torch.cumsum(sizes, dim=0).tolist()[:-1]
      )
      input = []
      for size in sizes:
          input += [torch.randperm(V)[:size]]
  
      input = torch.cat(input, dim=0)
  
      def chunk_method(x):
          input, offsets = x
          for i_offset, offset in enumerate(offsets):
              if i_offset < len(offsets) - 1:
                  next_offset = offsets[i_offset + 1]
              else:
                  next_offset = len(input)
              yield (input[offset:next_offset], torch.LongTensor([0]))
  
      self.run_test(
          (input, offsets),
          emb,
          chunk_method=chunk_method,
          ew_compatible=False,
      )
  
file_path: opacus/opacus/tests/grad_samples/embedding_bag_test.py
imports: import hypothesis.strategies
  import torch
  import torch.nn
  from hypothesis import given, settings
  from .common import GradSampleHooks_test
label_desc: Use image node_id 1 for calling the Embedding_bag_test obj's underlying member method code with example usage: obj.test_input_across_dims(N, sz, V, D, mode) without return types