code: def test_fm_shape(self) -> None:
      batch_size = 3
      # the input embedding are in torch.Tensor of [batch_size, num_embeddings, embedding_dim]
      input_embeddings = [
          torch.randn(batch_size, 2, 64),
          torch.randn(batch_size, 2, 32),
          torch.randn(batch_size, 3, 100),
          torch.randn(batch_size, 5, 120),
      ]
  
      fm = FactorizationMachine()
  
      fm_output = fm(input_embeddings)
  
      self.assertEqual(list(fm_output.shape), [batch_size, 1])
  
file_path: torchrec/torchrec/modules/tests/test_deepfm.py
imports: import unittest
  import torch
  from torchrec.fx import Tracer
  from torchrec.modules.deepfm import DeepFM, FactorizationMachine
label_desc: Use image node_id 1 for calling the TestFM obj's underlying member method code with example usage: obj.test_fm_shape() without return types