code: def test_grid_cluster(test, dtype, device):
      if dtype == torch.bfloat16 and device == torch.device("cuda:0"):
          return
  
      pos = tensor(test["pos"], dtype, device)
      size = tensor(test["size"], dtype, device)
      start = tensor(test.get("start"), dtype, device)
      end = tensor(test.get("end"), dtype, device)
  
      cluster = grid_cluster(pos, size, start, end)
      assert cluster.tolist() == test["cluster"]
  
      jit = torch.jit.script(grid_cluster)
      assert torch.equal(jit(pos, size, start, end), cluster)
  
imports: from itertools import product
  import pytest
  import torch
  from torch_cluster import grid_cluster
  from torch_cluster.testing import devices, dtypes, tensor
label_desc: Use image node_id 1 for calling a global function with example usage: test_grid_cluster(test, dtype, device) without return types