code: def test_gpu_cached_feature():
      a = torch.tensor([[1, 2, 3], [4, 5, 6]]).to("cuda").float()
      b = (
          torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
          .to("cuda")
          .float()
      )
  
      feat_store_a = gb.GPUCachedFeature(gb.TorchBasedFeature(a), 2)
      feat_store_b = gb.GPUCachedFeature(gb.TorchBasedFeature(b), 1)
  
      # Test read the entire feature.
      assert torch.equal(feat_store_a.read(), a)
      assert torch.equal(feat_store_b.read(), b)
  
      # Test read with ids.
      assert torch.equal(
          feat_store_a.read(torch.tensor([0]).to("cuda")),
          torch.tensor([[1.0, 2.0, 3.0]]).to("cuda"),
      )
      assert torch.equal(
          feat_store_b.read(torch.tensor([1, 1]).to("cuda")),
          torch.tensor(
              [[[4.0, 5.0], [6.0, 7.0]], [[4.0, 5.0], [6.0, 7.0]]]
          ).to("cuda"),
      )
  
      # Test get the size of the entire feature with ids.
      assert feat_store_a.size() == torch.Size([3])
      assert feat_store_b.size() == torch.Size([2, 2])
  
      # Test update the entire feature.
      feat_store_a.update(
          torch.tensor([[0.0, 1.0, 2.0], [3.0, 5.0, 2.0]]).to("cuda")
      )
      assert torch.equal(
          feat_store_a.read(),
          torch.tensor([[0.0, 1.0, 2.0], [3.0, 5.0, 2.0]]).to("cuda"),
      )
  
      # Test update with ids.
      feat_store_a.update(
          torch.tensor([[2.0, 0.0, 1.0]]).to("cuda"),
          torch.tensor([0]).to("cuda"),
      )
      assert torch.equal(
          feat_store_a.read(),
          torch.tensor([[2.0, 0.0, 1.0], [3.0, 5.0, 2.0]]).to("cuda"),
      )
  
imports: import unittest
  import backend
  import torch
  from dgl import graphbolt
label_desc: Use image node_id 1 for calling a global function with example usage: test_gpu_cached_feature() without return types