Allow FBGEMM_TBE_BOUNDS_CHECK_MODE to take effect when using mode 4,5,6 (#3838)

Summary:
X-link: facebookresearch/FBGEMM#925


This diff allows to set V2 check bound mode via env var `FBGEMM_TBE_BOUNDS_CHECK_MODE`, by setting its value to 4 (V2_IGNORE), 5 (V2_WARNING), 6 (V2_FATAL)

Previously, we can only get v2 check bound mode if bounds_check_mode is set to mode prefixed with V2.

Reviewed By: sryap

Differential Revision: D71344486

Author: yuhuishi-convect

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -687,6 +687,10 @@ def __init__(  # noqa C901
         self.is_nobag: bool = self.pooling_mode == PoolingMode.NONE
         # If environment variable is set, it overwrites the default bounds check mode.
         self.bounds_check_version: int = 1
+        self.bounds_check_mode_int: int = int(
+            os.environ.get("FBGEMM_TBE_BOUNDS_CHECK_MODE", bounds_check_mode.value)
+        )
+        bounds_check_mode = BoundsCheckMode(self.bounds_check_mode_int)
         if bounds_check_mode.name.startswith("V2_"):
             self.bounds_check_version = 2
             if bounds_check_mode == BoundsCheckMode.V2_IGNORE:
@@ -700,9 +704,6 @@ def __init__(  # noqa C901
                     f"Did not recognize V2 bounds check mode: {bounds_check_mode}"
                 )
 
-        self.bounds_check_mode_int: int = int(
-            os.environ.get("FBGEMM_TBE_BOUNDS_CHECK_MODE", bounds_check_mode.value)
-        )
         self.weights_precision = weights_precision
 
         if torch.cuda.is_available() and torch.version.hip: