force single gpu in inference.ipynb

2018-05-05 17:29:09 -07:00 · 2018-05-05 17:29:09 -07:00 · c67ca6531e
parent 78d5150d83
commit c67ca6531e
1 changed files with 4 additions and 1 deletions
--- a/inference.ipynb
+++ b/inference.ipynb
@ -98,8 +98,11 @@
   "source": [
    "checkpoint_path = \"/home/scratch.adlr-gcf/audio_denoising/runs/TTS-Tacotron2-LJS-MSE-DRC-NoMaskPadding-Unsorted-Distributed-22khz/checkpoint_15500\"\n",
    "model = load_model(hparams)\n",
+    "try:\n",
+    "    model = model.module\n",
+    "except:\n",
+    "    pass\n"
    "model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(checkpoint_path)['state_dict'].items()})\n",
-    "model = model.module\n",
    "_ = model.eval()"
   ]
  },