auto cuda

model.py

@@ -1,7 +1,8 @@
 from cvae import CVAE
 import torch
 from typing import Sequence
-import re
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 instruments = ['bass_acoustic', 'brass_acoustic', 'flute_acoustic', 'guitar_acoustic', 'keyboard_acoustic', 'mallet_acoustic', 'organ_acoustic', 'reed_acoustic', 'string_acoustic', 'synth_lead_acoustic', 'vocal_acoustic', 'bass_synthetic', 'brass_synthetic', 'flute_synthetic', 'guitar_synthetic', 'keyboard_synthetic', 'mallet_synthetic', 'organ_synthetic', 'reed_synthetic', 'string_synthetic', 'synth_lead_synthetic', 'vocal_synthetic', 'bass_electronic', 'brass_electronic', 'flute_electronic', 'guitar_electronic', 'keyboard_electronic', 'mallet_electronic', 'organ_electronic', 'reed_electronic', 'string_electronic', 'synth_lead_electronic', 'vocal_electronic']
 
@@ -13,7 +14,7 @@ model = CVAE.load_from_checkpoint(
     channels=[32, 64, 128, 256, 512],
     num_classes=len(instruments),
     learning_rate=1e-5
-)
+).to(device)
 
 def format(text):
     text = text.split(' ')[-1]
@@ -24,5 +25,5 @@ def choice_to_tensor(choice: Sequence[str]) -> torch.Tensor:
     return torch.tensor(instruments.index(choice))
 
 def generate(choice: Sequence[str], params: Sequence[int]=None):
-    noise = torch.tensor(params).unsqueeze(0).to('cuda') if params else torch.randn(1, 5).to('cuda')
-    return model.sample(eps=noise, c = choice_to_tensor(choice).to('cuda')).cpu().numpy()[0]
+    noise = torch.tensor(params).unsqueeze(0).to(device) if params else torch.randn(1, 5).to('cuda')
+    return model.sample(eps=noise, c = choice_to_tensor(choice).to(device)).cpu().numpy()[0]