Support MTIA in DenseTableBatchedEmbeddingBagsCodegen

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -2460,6 +2460,7 @@ def __init__(
         pooling_mode: PoolingMode = PoolingMode.SUM,
         use_cpu: bool = False,
         output_dtype: SparseType = SparseType.FP32,
+        use_mtia: bool = False,
     ) -> None:  # noqa C901  # tuple of (rows, dims,)
         super(DenseTableBatchedEmbeddingBagsCodegen, self).__init__()
 
@@ -2468,7 +2469,10 @@ def __init__(
         self.output_dtype: int = output_dtype.as_int()
         table_embedding_dtype = weights_precision.as_dtype()
 
-        self.use_cpu = use_cpu
+        self.use_cpu: bool = use_cpu
+        self.use_mtia: bool = use_mtia
+
+        assert not (use_cpu and use_mtia), "Cannot use CPU and MTIA at the same time"
 
         if self.use_cpu or self.pooling_mode == PoolingMode.NONE:
             assert output_dtype in [
@@ -2479,7 +2483,13 @@ def __init__(
 
         # pyre-fixme[8]: Attribute has type `device`; used as `Union[int, device]`.
         self.current_device: torch.device = (
-            torch.device("cpu") if self.use_cpu else torch.cuda.current_device()
+            torch.device("cpu")
+            if self.use_cpu
+            else (
+                torch.mtia.current_device()
+                if self.use_mtia
+                else torch.cuda.current_device()
+            )
         )
 
         self.embedding_specs = embedding_specs