diff --git a/multigen/pipes.py b/multigen/pipes.py
index df7d5e2..7e9a846 100755
--- a/multigen/pipes.py
+++ b/multigen/pipes.py
@@ -747,6 +747,8 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] =
                     self.pipe = self._class.from_pipe(self.pipe, controlnet=cnets)
             for cnet in cnets:
                 cnet.to(self.pipe.dtype)
+                if 'offload_device' not in args:
+                    cnet.to(self.pipe.device)
         else:
             # don't load anything, just reuse pipe
             super().__init__(model_id=model_id, pipe=pipe, **args)
diff --git a/tests/pipe_test.py b/tests/pipe_test.py
index 6690a69..8b6890e 100644
--- a/tests/pipe_test.py
+++ b/tests/pipe_test.py
@@ -162,9 +162,9 @@ def test_maskedimg2img_basic(self):
         scheduler = "EulerAncestralDiscreteScheduler"
         seed = 49045438434843
         blur = 48
-        param_3_3 = dict(original_image=img, image_painted=img_paint, strength=0.96,
+        param_3_3 = dict(image=img, image_painted=img_paint, strength=0.96,
                scheduler=scheduler, clip_skip=0, blur=blur, blur_compose=3, steps=5, guidance_scale=3.3)
-        param_7_6 = dict(original_image=img, image_painted=img_paint, strength=0.96,
+        param_7_6 = dict(image=img, image_painted=img_paint, strength=0.96,
                scheduler=scheduler, clip_skip=0, blur=blur, blur_compose=3, steps=5, guidance_scale=7.6)
         pipe.setup(**param_3_3)
         self.assertEqual(3.3, pipe.pipe_params['guidance_scale'])