code: def test_remove_diag(dtype, device):
      row, col = tensor(
          [[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device
      )
      value = tensor([1, 2, 3, 4], dtype, device)
      mat = SparseTensor(row=row, col=col, value=value)
      mat.fill_cache_()
  
      mat = mat.remove_diag()
      assert mat.storage.row().tolist() == [0, 1]
      assert mat.storage.col().tolist() == [1, 2]
      assert mat.storage.value().tolist() == [2, 3]
      assert mat.storage.num_cached_keys() == 2
      assert mat.storage.rowcount().tolist() == [1, 1, 0]
      assert mat.storage.colcount().tolist() == [0, 1, 1]
  
      mat = SparseTensor(row=row, col=col, value=value)
      mat.fill_cache_()
  
      mat = mat.remove_diag(k=1)
      assert mat.storage.row().tolist() == [0, 2]
      assert mat.storage.col().tolist() == [0, 2]
      assert mat.storage.value().tolist() == [1, 4]
      assert mat.storage.num_cached_keys() == 2
      assert mat.storage.rowcount().tolist() == [1, 0, 1]
      assert mat.storage.colcount().tolist() == [1, 0, 1]
  
file_path: pytorch_sparse/test/test_diag.py
imports: from itertools import product
  import pytest
  import torch
  from torch_sparse.tensor import SparseTensor
  from torch_sparse.testing import devices, dtypes, tensor
label_desc: Use image node_id 1 for calling a global function with example usage: test_remove_diag(dtype, device) without return types