Changing Backend Tensor initialization

Summary:
X-link: meta-pytorch/torchrec#3484

X-link: facebookresearch/FBGEMM#2066

**Context:**
Currently, we are enabling SSD optimizer offloading for the ssd tbe kernel

**In this diff:**
We retrieve the newly added parameters from the tbe config and pass it down to the tbe

Differential Revision: D85353134

Author: Raahul Kalyaan Jakka

fbgemm_gpu/src/ssd_split_embeddings_cache/embedding_rocksdb_wrapper.h

@@ -45,7 +45,9 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
       std::optional<at::Tensor> table_dims = std::nullopt,
       std::optional<at::Tensor> hash_size_cumsum = std::nullopt,
       int64_t flushing_block_size = 2000000000 /*2GB*/,
-      bool disable_random_init = false)
+      bool disable_random_init = false,
+      bool enable_optimizer_offloading = false,
+      int64_t optimizer_D = 0)
       : impl_(std::make_shared<ssd::EmbeddingRocksDB>(
             path,
             num_shards,
@@ -76,7 +78,9 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
             table_dims,
             hash_size_cumsum,
             flushing_block_size,
-            disable_random_init)) {}
+            disable_random_init,
+            enable_optimizer_offloading,
+            optimizer_D)) {}
 
   void set_cuda(
       at::Tensor indices,

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -809,7 +809,9 @@ static auto embedding_rocks_db_wrapper =
                 std::optional<at::Tensor>,
                 std::optional<at::Tensor>,
                 int64_t,
-                bool>(),
+                bool,
+                bool,
+                int64_t>(),
             "",
             {
                 torch::arg("path"),
@@ -842,6 +844,8 @@ static auto embedding_rocks_db_wrapper =
                 torch::arg("hash_size_cumsum") = std::nullopt,
                 torch::arg("flushing_block_size") = 2000000000 /* 2GB */,
                 torch::arg("disable_random_init") = false,
+                torch::arg("enable_optimizer_offloading") = false,
+                torch::arg("optimizer_D") = 0,
             })
         .def(
             "set_cuda",

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h

@@ -1378,6 +1378,7 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
   std::vector<std::string> db_paths_;
 
   bool disable_random_init_;
+  bool enable_optimizer_offloading;
 }; // class EmbeddingRocksDB
 
 /// @ingroup embedding-ssd