From 4280c512ae402d4e6035c65443a766212e87e65b Mon Sep 17 00:00:00 2001 From: Derek Hollowood Date: Fri, 10 Oct 2025 15:26:20 -0700 Subject: [PATCH] Fix for Missing _hash_zch_bucket Weight Issue Differential Revision: D84374871 --- torchrec/modules/hash_mc_modules.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchrec/modules/hash_mc_modules.py b/torchrec/modules/hash_mc_modules.py index a29dc3111..f7d566e99 100644 --- a/torchrec/modules/hash_mc_modules.py +++ b/torchrec/modules/hash_mc_modules.py @@ -305,7 +305,9 @@ 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]), persistent=False + ) # Do not need to store in buffer since this is created and consumed # at each step https://fburl.com/code/axzimmbx self._evicted_indices = []