made the files run on gpu

app.py

@@ -29,7 +29,7 @@ audio_path = 'models/pitch_to_audio/'
 # audio_path = '/network/scratch/n/nithya.shikarpur/checkpoints/pitch-diffusion/corrected-attention-v3/4835364'
 # db_path_audio = '/home/mila/n/nithya.shikarpur/scratch/pitch-diffusion/data/merged_data-finalest/cached-audio-pitch-16k'
 
-device = 'cpu'
+device = 'cuda'
 
 global_ind = -1
 global_audios = np.array([0.0])

src/generate_utils.py

@@ -61,7 +61,7 @@ def load_processed_pitch(pitch,
 def load_pitch_model(config, ckpt, qt = None, prime_file=None, device='cuda'):
     gin.parse_config_file(config)
     model = UNet()
-    model.load_state_dict(torch.load(ckpt, map_location='cpu')['state_dict'])  
+    model.load_state_dict(torch.load(ckpt, map_location='cuda')['state_dict'])  
     model.to(device)
     if qt is not None:
         qt = joblib.load(qt)
@@ -80,7 +80,7 @@ def load_pitch_model(config, ckpt, qt = None, prime_file=None, device='cuda'):
 def load_audio_model(config, ckpt, qt = None, device='cuda'):
     gin.parse_config_file(config)
     model = UNetPitchConditioned() # there are no gin parameters for some reason
-    model.load_state_dict(torch.load(ckpt, map_location='cpu')['state_dict'])  
+    model.load_state_dict(torch.load(ckpt, map_location='cuda')['state_dict'])  
     model.to(device)
     if qt is not None:
         qt = joblib.load(qt)