bofenghuang commited on
Commit
0ebb542
β€’
1 Parent(s): e59ead3
Files changed (1) hide show
  1. run_demo_layout.py +2 -1
run_demo_layout.py CHANGED
@@ -90,7 +90,8 @@ def _print_cuda_memory_info():
90
 
91
  def print_memory_info():
92
  _print_memory_info()
93
- _print_cuda_memory_info()
 
94
 
95
 
96
  def maybe_load_cached_pipeline(model_name):
 
90
 
91
  def print_memory_info():
92
  _print_memory_info()
93
+ if torch.cuda.is_available():
94
+ _print_cuda_memory_info()
95
 
96
 
97
  def maybe_load_cached_pipeline(model_name):