code: def remove_diag(src: SparseTensor, k: int = 0) -> SparseTensor:
      row, col, value = src.coo()
      inv_mask = row != col if k == 0 else row != (col - k)
      new_row, new_col = row[inv_mask], col[inv_mask]
  
      if value is not None:
          value = value[inv_mask]
  
      rowcount = src.storage._rowcount
      colcount = src.storage._colcount
      if rowcount is not None or colcount is not None:
          mask = ~inv_mask
          if rowcount is not None:
              rowcount = rowcount.clone()
              rowcount[row[mask]] -= 1
          if colcount is not None:
              colcount = colcount.clone()
              colcount[col[mask]] -= 1
  
      storage = SparseStorage(
          row=new_row,
          rowptr=None,
          col=new_col,
          value=value,
          sparse_sizes=src.sparse_sizes(),
          rowcount=rowcount,
          colptr=None,
          colcount=colcount,
          csr2csc=None,
          csc2csr=None,
          is_sorted=True,
      )
      return src.from_storage(storage)
  
file_path: pytorch_sparse/torch_sparse/diag.py
imports: from typing import Optional
  import torch
  from torch import Tensor
  from torch_sparse.storage import SparseStorage
  from torch_sparse.tensor import SparseTensor
label_desc: Use image node_id 1 for calling a global function with example usage: remove_diag(src, k) and returns: src