Update learned_round_optimizer.py

Xilinx · Dec 5, 2024 · 7cef5ce · 7cef5ce
1 parent 0f0dd43
commit 7cef5ce
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/brevitas_examples/common/learned_round/learned_round_optimizer.py b/src/brevitas_examples/common/learned_round/learned_round_optimizer.py
@@ -722,7 +722,8 @@ def skip_full_execution(self, block, next_block, block_forward, cache):
 
         # Finally (!), we compute the quantized input of the next block
         block.eval()
-        block.cuda()
+        if torch.cuda.is_available():
+          block.cuda()
         next_quant_input = []
         pbar = tqdm(range(len(cache)), desc='', leave=False)
         with torch.no_grad():
@@ -731,8 +732,8 @@ def skip_full_execution(self, block, next_block, block_forward, cache):
                 out = block_forward(block, (args, kwargs))
                 out = send_to_device(out, 'cpu')
                 next_quant_input.append((out,))
+        pbar.close()
         cache['args'] = next_quant_input
         block.cpu()
-        pbar.close()
 
         return cache