diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py
index 52dd4afd6b21..e97bdb352b22 100644
--- a/tests/pipelines/test_pipelines_common.py
+++ b/tests/pipelines/test_pipelines_common.py
@@ -698,11 +698,13 @@ def _test_xformers_attention_forwardGenerator_pass(
         pipe.set_progress_bar_config(disable=None)
 
         inputs = self.get_dummy_inputs(torch_device)
-        output_without_offload = pipe(**inputs)[0].cpu()
+        output_without_offload = pipe(**inputs)[0]
+        output_without_offload.cpu() if torch.is_tensor(output_without_offload) else output_without_offload
 
         pipe.enable_xformers_memory_efficient_attention()
         inputs = self.get_dummy_inputs(torch_device)
-        output_with_offload = pipe(**inputs)[0].cpu()
+        output_with_offload = pipe(**inputs)[0]
+        output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload
 
         if test_max_difference:
             max_diff = np.abs(output_with_offload - output_without_offload).max()