add cache sets assertion, and add logs for training and inference SSD…

… TBE (#2689)

Summary:
Pull Request resolved: #2689

Previous diffs added loggings and assertion, but they were added inconsistently, and to the wrong TBEs.

Reviewed By: sryap

Differential Revision: D58217853

fbshipit-source-id: 63663df6d46bde42b25305d5f42e95e1b06a9ccd

fbgemm_gpu/fbgemm_gpu/ssd_split_table_batched_embeddings_ops.py

@@ -177,6 +177,7 @@ def __init__(
             ),
         )
 
+        assert cache_sets > 0
         element_size = 4
         cache_size = cache_sets * ASSOC * element_size * self.max_D
         logging.info(
@@ -885,10 +886,12 @@ def max_ty_D(ty: SparseType) -> int:
         logging.info(
             f"Using cache for SSD with admission algorithm "
             f"{CacheAlgorithm.LRU}, {cache_sets} sets, stored on {'DEVICE' if ssd_cache_location is EmbeddingLocation.DEVICE else 'MANAGED'} with {ssd_shards} shards, "
+            f"SSD storage directory: {ssd_storage_directory}, "
             f"Memtable Flush Period: {ssd_memtable_flush_period}, "
             f"Memtable Flush Offset: {ssd_memtable_flush_offset}, "
             f"Desired L0 files per compaction: {ssd_l0_files_per_compact}, "
-            f"{cache_size / 1024.0 / 1024.0 / 1024.0 : .2f}GB"
+            f"{cache_size / 1024.0 / 1024.0 / 1024.0 : .2f}GB, "
+            f"output dtype: {output_dtype}"
         )
         self.register_buffer(
             "lxu_cache_state",