[Bugfix] Optimize CpuGpuBuffer initialization (#25447)

Signed-off-by: Naman Lalit <nl2688@nyu.edu>
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Naman Lalit 2025-09-26 17:07:36 -07:00 committed by yewentao256
parent c4b9864e22
commit e7cba8f6b1

View File

@ -117,7 +117,7 @@ class CpuGpuBuffer:
dtype=dtype,
device="cpu",
pin_memory=pin_memory)
self.gpu = self.cpu.to(device)
self.gpu = torch.zeros_like(self.cpu, device=device)
self.np: np.ndarray
# To keep type hints simple (avoiding generics and subclasses), we
# only conditionally create the numpy array attribute. This can cause