remove output dtype restriction in ssd tbe (#3524)

Summary: X-link: facebookresearch/FBGEMM#606 as title Differential Revision: D67563019
pytorch · Dec 22, 2024 · 05a667d · 05a667d
1 parent fe980ab
commit 05a667d
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py
@@ -290,7 +290,6 @@ def __init__(
         )
 
         assert weights_precision in (SparseType.FP32, SparseType.FP16)
-        assert output_dtype in (SparseType.FP32, SparseType.FP16)
         self.weights_precision = weights_precision
         self.output_dtype: int = output_dtype.as_int()