diff --git a/optimum/onnxruntime/modeling_diffusion.py b/optimum/onnxruntime/modeling_diffusion.py
index 996dbf2f49..a0a33d3cb7 100644
--- a/optimum/onnxruntime/modeling_diffusion.py
+++ b/optimum/onnxruntime/modeling_diffusion.py
@@ -879,7 +879,7 @@ class ORTDiffusionPipeline(ConfigMixin):
 
     @classmethod
     @validate_hf_hub_args
-    def from_pretrained(cls, pretrained_model_or_path, **kwargs):
+    def from_pretrained(cls, pretrained_model_or_path, **kwargs) -> ORTPipeline:
         load_config_kwargs = {
             "force_download": kwargs.get("force_download", False),
             "resume_download": kwargs.get("resume_download", None),
@@ -953,7 +953,7 @@ class ORTPipelineForTask(ConfigMixin):
     config_name = "model_index.json"
 
     @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, **kwargs):
+    def from_pretrained(cls, pretrained_model_or_path, **kwargs) -> ORTPipeline:
         load_config_kwargs = {
             "force_download": kwargs.get("force_download", False),
             "resume_download": kwargs.get("resume_download", None),
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
index 1a1f84c884..1da49a65a2 100644
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -407,19 +407,18 @@ def evaluation_loop(
 
 def np_to_pt(np_object, device):
     if isinstance(np_object, np.ndarray):
-        if np_object.ndim == 4:
-            return torch.from_numpy(np_object).permute(0, 3, 1, 2)
-        elif np_object.ndim == 3:
-            return torch.from_numpy(np_object).permute(2, 0, 1)
-        else:
-            return torch.from_numpy(np_object)
-    elif isinstance(np_object, list) and isinstance(np_object[0], np.ndarray):
-        return [np_to_pt(a, device) for a in np_object]
-    elif isinstance(np_object, dict) and isinstance(next(iter(np_object.values())), np.ndarray):
-        return {k: np_to_pt(v, device) for k, v in np_object.items()}
+        return torch.from_numpy(np_object)
     elif isinstance(np_object, np.random.RandomState):
         return torch.Generator(device=device).manual_seed(int(np_object.get_state()[1][0]))
-    elif isinstance(np_object, list) and isinstance(np_object[0], np.random.RandomState):
-        return [torch.Generator(device=device).manual_seed(int(a.get_state()[1][0])) for a in np_object]
+    elif isinstance(np_object, np.random.Generator):
+        return torch.Generator(device=device).manual_seed(int(np_object.bit_generator.state[1][0]))
+    elif isinstance(np_object, list) and isinstance(
+        np_object[0], (np.ndarray, np.random.RandomState, np.random.Generator)
+    ):
+        return [np_to_pt(a, device) for a in np_object]
+    elif isinstance(np_object, dict) and isinstance(
+        next(iter(np_object.values())), (np.ndarray, np.random.RandomState, np.random.Generator)
+    ):
+        return {k: np_to_pt(v, device) for k, v in np_object.items()}
     else:
         return np_object
diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py
index 233ac29513..cdcee7f613 100644
--- a/tests/onnxruntime/test_diffusion.py
+++ b/tests/onnxruntime/test_diffusion.py
@@ -62,7 +62,7 @@ def _generate_images(height=128, width=128, batch_size=1, channel=3, input_type=
             "/in_paint/overture-creations-5sI6fQgYIuo.png"
         ).resize((width, height))
     elif input_type == "np":
-        image = np.random.rand(height, width, channel)
+        image = np.random.rand(channel, height, width)
     elif input_type == "pt":
         image = torch.rand((channel, height, width))
 
@@ -461,10 +461,9 @@ def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str):
         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
 
         pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[test_name], provider=provider)
+        self.assertEqual(pipeline.device.type, "cuda")
+
         outputs = pipeline(**inputs).images
-        # Verify model devices
-        self.assertEqual(pipeline.device.type.lower(), "cuda")
-        # Verify model outptus
         self.assertIsInstance(outputs, np.ndarray)
         self.assertEqual(outputs.shape, (batch_size, height, width, 3))
 
@@ -650,9 +649,8 @@ def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str):
         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
 
         pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[test_name], provider=provider)
+        self.assertEqual(pipeline.device, "cuda")
+
         outputs = pipeline(**inputs).images
-        # Verify model devices
-        self.assertEqual(pipeline.device.type.lower(), "cuda")
-        # Verify model outptus
         self.assertIsInstance(outputs, np.ndarray)
         self.assertEqual(outputs.shape, (batch_size, height, width, 3))