Unverified Commit 5de8d9f1 authored by Maximilien de Bayser's avatar Maximilien de Bayser Committed by GitHub
Browse files

Remove extra tensor on CPU (#20693)


Signed-off-by: default avatarMax de Bayser <mbayser@br.ibm.com>
parent c1c8ca57
......@@ -234,10 +234,16 @@ class MinPLogitsProcessor(LogitsProcessor):
device="cpu",
pin_memory=pin_memory)
self.min_p_cpu = self.min_p_cpu_tensor.numpy()
# Pre-allocated device tensor
self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
dtype=torch.float32,
device=device)
self.use_double_tensor = torch.device("cpu") != torch.device(device)
if self.use_double_tensor:
# Pre-allocated device tensor
self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
dtype=torch.float32,
device=device)
else:
self.min_p_device = self.min_p_cpu_tensor
# Current slice of the device tensor
self.min_p: torch.Tensor = self.min_p_device[:0]
......@@ -284,7 +290,9 @@ class MinPLogitsProcessor(LogitsProcessor):
size = batch_update.batch_size
if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
self.min_p = self.min_p_device[:size]
self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
if self.use_double_tensor:
self.min_p.copy_(self.min_p_cpu_tensor[:size],
non_blocking=True)
self.min_p.unsqueeze_(1)
def apply(self, logits: torch.Tensor) -> torch.Tensor:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment