Update collect_diffusion_input_imagenet.py

remove setting CUDA_VISIBLE_DEVICES='7'
ziplab · Mar 12, 2024 · 6b190f9 · 6b190f9
1 parent 84a120a
commit 6b190f9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/quant_scripts/collect_diffusion_input_imagenet.py b/quant_scripts/collect_diffusion_input_imagenet.py
@@ -3,7 +3,7 @@
 sys.path.append(".")
 sys.path.append('./taming-transformers')
 import os
-os.environ['CUDA_VISIBLE_DEVICES'] = '7'
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 from taming.models import vqgan
 
 import torch
@@ -78,4 +78,4 @@ def get_model():
     import ldm.globalvar as globalvar   
     input_list = globalvar.getInputList()
     torch.save(input_list, 'imagenet_input_{}steps.pth'.format(ddim_steps))
-    sys.exit(0)
+    sys.exit(0)