Lint fix

apache · May 18, 2021 · 1808b3d · 1808b3d
1 parent adf0597
commit 1808b3d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/python/tvm/topi/sparse/utils.py b/python/tvm/topi/sparse/utils.py
@@ -188,6 +188,7 @@ def convert_model_dense_to_sparse(
 
 
 def sparse_sketch_rules():
+    """Return the sketch rules for sparse op"""
     sparse_sketch_rule_list = [
         auto_scheduler.PreloadCustomSketchRule(
             sparse_conv2d_meet_condition_func, sparse_conv2d_apply_func, "SparseConv2D"

diff --git a/tests/python/relay/test_sparse_conv2d_convert.py b/tests/python/relay/test_sparse_conv2d_convert.py
@@ -56,7 +56,9 @@ def test_bsr_sparse_conv2d_nchw():
 
     params = {
         "weight": tvm.nd.array(
-            np.array(random_bsr_matrix(128, 64, 8, 1, 0.1).todense()).reshape(128, 64, 1, 1)
+            np.array(random_bsr_matrix(128, 64, 8, 1, 0.1, "float32").todense()).reshape(
+                128, 64, 1, 1
+            )
         )
     }
 
@@ -81,7 +83,9 @@ def test_bsr_sparse_conv2d_nhwc():
 
     params = {
         "weight": tvm.nd.array(
-            np.array(random_bsr_matrix(128, 64, 8, 1, 0.1).todense()).T.reshape(1, 1, 64, 128)
+            np.array(random_bsr_matrix(128, 64, 8, 1, 0.1, "float32").todense()).T.reshape(
+                1, 1, 64, 128
+            )
         )
     }