Update extract_embeddings.py

NYUMedML · Aug 5, 2022 · 32c0551 · 32c0551
1 parent b58070c
commit 32c0551
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/feature_extraction/extract_embeddings.py b/feature_extraction/extract_embeddings.py
@@ -76,7 +76,7 @@ def load_pretrained(net, model_dir):
 feature_extractor = InceptionV4(num_classes=256)
 load_pretrained(feature_extractor, args.feature_extractor_dir)
 feature_extractor.to('cuda')
-featurizer = nn.DataParallel(model, device_ids=device_ids)
+feature_extractor = nn.DataParallel(feature_extractor, device_ids=device_ids)
 
 subtype_model = InceptionV4(num_classes=2).to('cuda')
 subtype_model.load_state_dict(torch.load(args.subtype_model_dir))