remove output dtype restriction in ssd tbe (#3524)

Summary: Pull Request resolved: #3524 X-link: facebookresearch/FBGEMM#606 as title Reviewed By: jiayulu Differential Revision: D67563019 fbshipit-source-id: 6c349c2bdc2b97f1455e52dc102fe674a75c375f
pytorch · Dec 27, 2024 · bf354fa · bf354fa
1 parent 7d41ee5
commit bf354fa
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()