Fix MPZCH model publish (#3464)

Summary:
Pull Request resolved: #3464

Updated _hash_zch_bucket as a 2d tensor to enable quantization for GR

Reviewed By: jeffreyjian, zlzhao1104

Differential Revision: D84944806

fbshipit-source-id: 8d46ba7741e88cd78d1ff5248f7bb98cbf975eec

Author: chaodengusc

torchrec/modules/hash_mc_modules.py

@@ -305,7 +305,7 @@ def __init__(
 
         self._max_probe = max_probe
         self._buckets = total_num_buckets
-        self.register_buffer("_hash_zch_bucket", torch.tensor([total_num_buckets]))
+        self.register_buffer("_hash_zch_bucket", torch.tensor([[total_num_buckets]]))
         # Do not need to store in buffer since this is created and consumed
         # at each step https://fburl.com/code/axzimmbx
         self._evicted_indices = []