diff --git a/fbgemm_gpu/test/tbe/training/merge_vbe_test.py b/fbgemm_gpu/test/tbe/training/merge_vbe_test.py
index de75145db3..87e7f71be2 100755
--- a/fbgemm_gpu/test/tbe/training/merge_vbe_test.py
+++ b/fbgemm_gpu/test/tbe/training/merge_vbe_test.py
@@ -9,6 +9,7 @@
 
 # pyre-ignore-all-errors[56]
 
+from __future__ import annotations
 
 import tempfile
 import unittest
@@ -84,7 +85,7 @@ def __init__(
         weights_precision: SparseType,
         feature_table_map: Optional[List[int]],
         is_ssd: bool = False,
-    ):
+    ) -> None:
         self.T = (
             len(embedding_specs)
             if feature_table_map is None
@@ -307,7 +308,7 @@ def execute_merge_vbe(
         features_list: Optional[list[int]] = None,
         dims_list: Optional[list[int]] = None,
         Bs_feature_rank_list: Optional[List[List[List[int]]]] = None,
-    ):
+    ) -> None:
         """
         Execute merge VBE
 
@@ -505,9 +506,9 @@ def test_merge_VBE_SSD(  # noqa C901
     @settings(**common_settings)
     def test_merge_VBE_nonSSD(  # noqa C901
         self,
-        optimizer,
-        weights_precision,
-        output_dtype,
+        optimizer: OptimType,
+        weights_precision: SparseType,
+        output_dtype: SparseType,
         **kwargs: Any,
     ) -> None:
         """
@@ -532,7 +533,7 @@ def test_merge_VBE_nonSSD(  # noqa C901
     @settings(**common_settings)
     def test_merge_VBE_simple(  # noqa C901
         self,
-        num_ssd,
+        num_ssd: int,
     ) -> None:
         """
         Test merge VBE simple case