def METHOD_NAME(self, tensor_shape, dtype, name):
    required_len = reduce(operator.mul, tensor_shape, 1)
    if (
        self.buffer.get((name, dtype), None) is None
        or self.buffer[(name, dtype)].numel() < required_len
    ):
        self.buffer[(name, dtype)] = torch.empty(
            required_len, dtype=dtype, device=torch.cuda.current_device(), requires_grad=False
        )
    return self.buffer[(name, dtype)][0:required_len].view(*tensor_shape)