fix device mismatch wavegrad training

this should fixe the device mismatch as seen here https://github.com/mozilla/TTS/issues/622#issue-789802916

TTS/bin/train_vocoder_wavegrad.py

@@ -344,6 +344,10 @@ def main(args):  # pylint: disable=redefined-outer-name
 
     # setup criterion
     criterion = torch.nn.L1Loss().cuda()
+    
+    if use_cuda:
+        model.cuda()
+        criterion.cuda()
 
     if args.restore_path:
         checkpoint = torch.load(args.restore_path, map_location='cpu')
@@ -378,10 +382,6 @@ def main(args):  # pylint: disable=redefined-outer-name
     else:
         args.restore_step = 0
 
-    if use_cuda:
-        model.cuda()
-        criterion.cuda()
-
     # DISTRUBUTED
     if num_gpus > 1:
         model = DDP_th(model, device_ids=[args.rank])