Make novel_predict actually use the GPU it asks for

CHANGELOG.md

@@ -32,8 +32,6 @@
 
 ## BUG FIXES
 
-* `cellmapper_scvi`: Bump the base image from `openproblems/base_pytorch_nvidia:1.0.0` to `:1`. (PR #53).
-
 * `process_dataset`: Fall back to holding out a quarter of the batches when the dataset has no `obs["is_train"]`, rather than silently producing four empty h5ads. `obs["is_train"]` carries the NeurIPS 2021 competition split and stays optional; `obs["cell_type"]` is now declared and required (PR #28).
 
 * Fix the component paths, build paths and `rename_keys` separator in the helper scripts, which prevented `scripts/create_datasets/test_resources.sh` and both `run_test.sh` scripts from running at all (PR #22).
@@ -58,6 +56,10 @@
 
 * `simple_mlp`: Drop the `input_transform` key from the `simple_mlp_predict` call, which is not an argument of that component (PR #44).
 
+* `cellmapper_scvi`: Bump the base image from `openproblems/base_pytorch_nvidia:1.0.0` to `:1`. (PR #53).
+
+* `novel_predict`: Move the model and the input batch onto the selected device. It picked `cuda:0` when a GPU was present but never used it, so the component requested a GPU node and ran inference on CPU (PR #54).
+
 # task_predict_modality 0.1.1
 
 ## NEW FUNCTIONALITY

src/methods/novel/novel_predict/script.py

@@ -95,10 +95,13 @@
 dataloader_test = DataLoader(dataset_test, 32, shuffle = False, num_workers = 4)
 
 outputs = []
+# the weights are loaded with map_location='cpu', so move the model onto the device
+# we selected above -- otherwise this component requests a GPU and never uses it
+model = model.to(device)
 model.eval()
 with torch.no_grad():
     for x in dataloader_test:
-        output = model(x.float())
+        output = model(x.float().to(device))
         outputs.append(output.detach().cpu().numpy())
 
 outputs = np.concatenate(outputs)