From 33010f44b1a451bb6e7e9fcb4a2c674854db77dc Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 27 Mar 2024 12:26:49 +0200
Subject: [PATCH] Attempt to fix ensembling on MPS

---
 marigold/util/ensemble.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/marigold/util/ensemble.py b/marigold/util/ensemble.py
index 85c8f77..d34b153 100644
--- a/marigold/util/ensemble.py
+++ b/marigold/util/ensemble.py
@@ -79,8 +79,13 @@ def ensemble_depths(input_images, regularizer_strength=0.02, max_iter=2, tol=1e-
     t = x[int(l/2):]
     
     # Prediction
-    s = torch.from_numpy(s).to(device)
-    t = torch.from_numpy(t).to(device)
+    try:
+        s = torch.from_numpy(s).to(device)
+        t = torch.from_numpy(t).to(device)
+    except:
+        s = torch.from_numpy(s.astype(np.float32)).to(device)
+        t = torch.from_numpy(t.astype(np.float32)).to(device)
+
     transformed_arrays = original_input * s.view(-1, 1, 1) + t.view(-1, 1, 1)
     if 'mean' == reduction:
         aligned_images = torch.mean(transformed_arrays, dim=0)