diff --git a/README.md b/README.md
index 5222bfd..17076e0 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ import numpy as np
 from torch.utils.data import DataLoader
 from torchvision.transforms import v2
 from torchvision.transforms import InterpolationMode
-from github.datasets.h5py_dataset import H5PYDataset
+from phoenix.datasets.h5py_dataset import H5PYDataset
 
 gene_path = './xenium_human_multi.npy'
 gene_list = list(np.load(gene_path))
@@ -122,7 +122,7 @@ print("Output:", output.size())
 
 To predict gene expression from histology images use
 ```python
-from github.helpers.inference import FlowPipeline
+from phoenix.helpers.inference import FlowPipeline
 
 pipeline = FlowPipeline(
     model=flow_model,
diff --git a/phoenix/helpers/inference.py b/phoenix/helpers/inference.py
index 60c14b7..0532e79 100644
--- a/phoenix/helpers/inference.py
+++ b/phoenix/helpers/inference.py
@@ -63,9 +63,9 @@ def __call__(self, gene_list: list, dataloader: DataLoader):
 
         pred_list, coords_list = [], []
         for batch in tqdm(dataloader, desc='Flow sampling'):
-            image, coords = batch[0].cuda(), batch[1]
+            image, coords = batch[0].to(device), batch[1]
             feats = self.model.vision_forward(image)
-            noise = torch.randn(image.size(0), len(gene_list), 1).cuda()
+            noise = torch.randn(image.size(0), len(gene_list), 1, device=device)
 
             gex_pred = run_flow(
                 flow_model=self.model,