jhj0517 commited on
Commit
4f7f3bd
1 Parent(s): b2c613c
Files changed (1) hide show
  1. modules/base_interface.py +3 -2
modules/base_interface.py CHANGED
@@ -9,8 +9,9 @@ class BaseInterface:
9
 
10
  @staticmethod
11
  def release_cuda_memory():
12
- torch.cuda.empty_cache()
13
- torch.cuda.reset_max_memory_allocated()
 
14
 
15
  @staticmethod
16
  def remove_input_files(file_paths: List[str]):
 
9
 
10
  @staticmethod
11
  def release_cuda_memory():
12
+ if torch.cuda.is_available():
13
+ torch.cuda.empty_cache()
14
+ torch.cuda.reset_max_memory_allocated()
15
 
16
  @staticmethod
17
  def remove_input_files(file_paths: List[str]):