refactor to_sparse_semi_structure

ROCm · Oct 3, 2024 · 36f5443 · 36f5443
1 parent b90a8d3
commit 36f5443
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu b/aten/src/ATen/native/sparse/cuda/SparseSemiStructuredOps.cu
@@ -908,14 +908,14 @@ static void reorder_meta(cutlass::TensorRef<Element, LayoutDest> dest,
 
 std::tuple<Tensor, Tensor>
 _to_sparse_semi_structured(const Tensor& dense) {
+  TORCH_CHECK(dense.dim() == 2,
+              __func__, " : Expected dense argument to be 2D tensor, got ",
+              dense.dim(), " dims");
 #if defined(_MSC_VER) || (defined(CUDA_VERSION) && CUDA_VERSION < 11080)
   AT_ERROR(__func__, " : CUTLASS not supported");
   return std::make_tuple(Tensor{}, Tensor{});
 #elif defined(USE_ROCM)
   // Check dimensions of the dense matrix.
-  TORCH_CHECK(dense.dim() == 2,
-              __func__, " : Expected dense argument to be 2D tensor, got ",
-              dense.dim(), " dims");
 
   // Generate sparse tensor using cuSPARSELt compression
   auto sparse = torch._cslt_compress(dense);