code: def forward(
      self, features: "KeyedJaggedTensor"
  ) -> "KeyedJaggedTensor":
      weights_list = []
      for key, boundaries in self.boundaries_dict.items():
          jt = features[key]
          bucketized = torch.bucketize(jt.weights(), boundaries)
          # doesn't super matter I guess
          # hashed = torch.ops.fb.index_hash(bucketized, seed=0, modulo=len(boundaries))
          hashed = bucketized
          weights = torch.gather(
              self.bucket_w[key], dim=0, index=hashed
          )
          weights_list.append(weights)
      return KeyedJaggedTensor(
          keys=features.keys(),
          values=features.values(),
          weights=torch.cat(weights_list),
          lengths=features.lengths(),
          offsets=features.offsets(),
          stride=features.stride(),
          length_per_key=features.length_per_key(),
      )
  
file_path: pytorch/test/dynamo/test_torchrec.py
imports: import sys
  import unittest
  from typing import Dict, List
  import torch
  import torch._dynamo.config
  import torch._dynamo.test_case
  from torch import nn
  from torch._dynamo.test_case import TestCase
  from torch._dynamo.testing import CompileCounter
  from torch.testing._internal.common_utils import NoTest
label_desc: Use image node_id 2 for calling the BucketizeMod obj's underlying member method code with example usage: obj.forward(features) and returns: KeyedJaggedTensor