Add CacheState type to fbgemm_gpu (pytorch#591)

Summary: Pull Request resolved: pytorch#591 Reviewed By: jianyuh, xush6528 Differential Revision: D27607134 fbshipit-source-id: 7763a3743b3cb3f95154def2b84a46ba363aefd8
joebos · Apr 7, 2021 · 4b5e997 · 4b5e997
1 parent b01c349
commit 4b5e997
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py
@@ -936,8 +936,7 @@ def _apply_split(self, split: SplitState, prefix: str, dtype: torch.dtype, enfor
 
     def _apply_cache_state(
         self,
-        # pyre-fixme[2]: Parameter must be annotated.
-        cache_state,
+        cache_state: CacheState,
         cache_algorithm: CacheAlgorithm,
         cache_load_factor: float,
         cache_sets: int,