[Bugfix] Optimize CpuGpuBuffer initialization (#25447)
Signed-off-by: Naman Lalit <nl2688@nyu.edu>
This commit is contained in:
@@ -117,7 +117,7 @@ class CpuGpuBuffer:
|
|||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
device="cpu",
|
device="cpu",
|
||||||
pin_memory=pin_memory)
|
pin_memory=pin_memory)
|
||||||
self.gpu = self.cpu.to(device)
|
self.gpu = torch.zeros_like(self.cpu, device=device)
|
||||||
self.np: np.ndarray
|
self.np: np.ndarray
|
||||||
# To keep type hints simple (avoiding generics and subclasses), we
|
# To keep type hints simple (avoiding generics and subclasses), we
|
||||||
# only conditionally create the numpy array attribute. This can cause
|
# only conditionally create the numpy array attribute. This can cause
|
||||||
|
|||||||
Reference in New Issue
Block a user