fix #6799

Signed-off-by: KumoLiu <[email protected]>
Project-MONAI · Jul 31, 2023 · 3d3e267 · 3d3e267
1 parent 5feb353
commit 3d3e267
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/tests/test_selfattention.py b/tests/test_selfattention.py
@@ -62,7 +62,7 @@ def test_access_attn_matrix(self):
         # be  not able to access the matrix
         no_matrix_acess_blk = SABlock(hidden_size=hidden_size, num_heads=num_heads, dropout_rate=dropout_rate)
         no_matrix_acess_blk(torch.randn(input_shape))
-        assert type(no_matrix_acess_blk.att_mat) == torch.Tensor
+        assert isinstance(no_matrix_acess_blk.att_mat, torch.Tensor)
         # no of elements is zero
         assert no_matrix_acess_blk.att_mat.nelement() == 0
 

diff --git a/tests/test_transformerblock.py b/tests/test_transformerblock.py
@@ -66,7 +66,7 @@ def test_access_attn_matrix(self):
             hidden_size=hidden_size, mlp_dim=mlp_dim, num_heads=num_heads, dropout_rate=dropout_rate
         )
         no_matrix_acess_blk(torch.randn(input_shape))
-        assert type(no_matrix_acess_blk.attn.att_mat) == torch.Tensor
+        assert isinstance(no_matrix_acess_blk.attn.att_mat, torch.Tensor)
         # no of elements is zero
         assert no_matrix_acess_blk.attn.att_mat.nelement() == 0
 

diff --git a/tests/test_vit.py b/tests/test_vit.py
@@ -160,7 +160,7 @@ def test_access_attn_matrix(self):
         # no data in the matrix
         no_matrix_acess_blk = ViT(in_channels=in_channels, img_size=img_size, patch_size=patch_size)
         no_matrix_acess_blk(torch.randn(in_shape))
-        assert type(no_matrix_acess_blk.blocks[0].attn.att_mat) == torch.Tensor
+        assert isinstance(no_matrix_acess_blk.blocks[0].attn.att_mat, torch.Tensor)
         # no of elements is zero
         assert no_matrix_acess_blk.blocks[0].attn.att_mat.nelement() == 0