Fix (gptq): pin_memory only available with CUDA

Xilinx · Aug 27, 2024 · 1f6432b · 1f6432b
1 parent 059fdc1
commit 1f6432b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/brevitas/graph/gptq.py b/src/brevitas/graph/gptq.py
@@ -133,11 +133,11 @@ def __init__(
         self.H = torch.zeros((self.groups, self.columns, self.columns),
                              device='cpu',
                              dtype=torch.float32,
-                             pin_memory=True)
+                             pin_memory=torch.cuda.is_available())
         self.B = torch.zeros((self.groups, self.columns, self.columns),
                              device='cpu',
                              dtype=torch.float32,
-                             pin_memory=True)
+                             pin_memory=torch.cuda.is_available())
         self.nsamples = 0
 
         assert torch_version >= version.parse('1.10'), "GPTQ requires torch 1.10 or higher"