diff --git a/app.py b/app.py
index 4dff0b972..9172df51a 100644
--- a/app.py
+++ b/app.py
@@ -631,7 +631,9 @@ def _do_predictions(gen_type, texts, melodies, sample, trim_start, trim_end, dur
     if UNLOAD_MODEL:
         MODEL = None
     torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
+    # fix not install cuda driver
+    if torch.cuda.is_available():
+        torch.cuda.ipc_collect()
     return res, res_audio, res_backup, input_length
 
 
@@ -1832,4 +1834,4 @@ def ui_batched(launch_kwargs):
         USE_DIFFUSION = False
         ui_batched(launch_kwargs)
     else:
-        ui_full(launch_kwargs)
\ No newline at end of file
+        ui_full(launch_kwargs)
diff --git a/audiocraft/models/audiogen.py b/audiocraft/models/audiogen.py
index 6adefb974..d0fcaa3f3 100644
--- a/audiocraft/models/audiogen.py
+++ b/audiocraft/models/audiogen.py
@@ -71,16 +71,11 @@ def audio_channels(self) -> int:
         return self.compression_model.channels
 
     @staticmethod
-    def get_pretrained(name: str = 'facebook/audiogen-medium', device=None):
+    def get_pretrained(name: str = 'facebook/audiogen-medium', device=('cuda' if torch.cuda.is_available() else 'cpu')):
         """Return pretrained model, we provide a single model for now:
         - facebook/audiogen-medium (1.5B), text to sound,
           # see: https://huggingface.co/facebook/audiogen-medium
         """
-        if device is None:
-            if torch.cuda.device_count():
-                device = 'cuda'
-            else:
-                device = 'cpu'
 
         if name == 'debug':
             # used only for unit tests
@@ -273,4 +268,4 @@ def _progress_callback(generated_tokens: int, tokens_to_generate: int):
     def to(self, device: str):
         self.compression_model.to(device)
         self.lm.to(device)
-        return self
\ No newline at end of file
+        return self
diff --git a/audiocraft/models/musicgen.py b/audiocraft/models/musicgen.py
index 1d4b2292e..5130767d5 100644
--- a/audiocraft/models/musicgen.py
+++ b/audiocraft/models/musicgen.py
@@ -85,7 +85,7 @@ def audio_channels(self) -> int:
         return self.compression_model.channels
 
     @staticmethod
-    def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=None):
+    def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=('cuda' if torch.cuda.is_available() else 'cpu')):
         """Return pretrained model, we provide four models:
         - facebook/musicgen-small (300M), text to music,
           # see: https://huggingface.co/facebook/musicgen-small
@@ -96,11 +96,6 @@ def get_pretrained(name: str = 'GrandaddyShmax/musicgen-melody', device=None):
         - facebook/musicgen-large (3.3B), text to music,
           # see: https://huggingface.co/facebook/musicgen-large
         """
-        if device is None:
-            if torch.cuda.device_count():
-                device = 'cuda'
-            else:
-                device = 'cpu'
 
         if name == 'debug':
             # used only for unit tests
@@ -406,4 +401,4 @@ def generate_audio(self, gen_tokens: torch.Tensor):
     def to(self, device: str):
         self.compression_model.to(device)
         self.lm.to(device)
-        return self
\ No newline at end of file
+        return self