code: def test_collapse():
      with torch.no_grad():  # remove a warning
          size = (5, 6)
          param = torch.rand(size, requires_grad=True)
          param.grad = torch.rand(size)
  
          bucket = GradBucket(300, param.dtype, param.device, -1)
          bucket.add_grad(param)
          bucket.shrink()
          bucket.collapse()
  
          assert bucket.buffer.numel() == 0
          assert param.grad is None
          bucket.rebuild()
  
          assert param.grad is not None
          torch.allclose(param.grad, torch.zeros(size))
  
imports: import pytest
  import torch
  from fairscale.nn.misc import GradBucket
label_desc: Use image node_id 4 for calling a global function with example usage: test_collapse() without return types