Update single node learning with pretrained fetched model

etzinis · Apr 16, 2021 · 71c151f · 71c151f
1 parent 1f68383
commit 71c151f
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/approx_ensembles/fedenhance/experiments/run_unsup_sep_enh_v2.py b/approx_ensembles/fedenhance/experiments/run_unsup_sep_enh_v2.py
@@ -116,6 +116,11 @@
 else:
     raise ValueError('Invalid model: {}.'.format(hparams['model_type']))
 
+# Load the pre-trained model if it is given.
+if hparams['warmup_checkpoint'] is not None:
+    model.load_state_dict(
+        torch.load(hparams['warmup_checkpoint']))
+
 numparams = 0
 for f in model.parameters():
     if f.requires_grad: