From 49503c2175a71c3d71b75d5702c4ff9e8b0342cb Mon Sep 17 00:00:00 2001
From: Jose <34888496+Jerry-Master@users.noreply.github.com>
Date: Fri, 1 Dec 2023 01:04:01 +0100
Subject: [PATCH] Mps fix (#2436)

* mps bugfix

* removed docs change

* Double quotes

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>
---
 nerfstudio/engine/trainer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
index fe4aa85ce..4cc296063 100644
--- a/nerfstudio/engine/trainer.py
+++ b/nerfstudio/engine/trainer.py
@@ -464,6 +464,7 @@ class Trainer:
 
         self.optimizers.zero_grad_all()
         cpu_or_cuda_str: str = self.device.split(":")[0]
+        cpu_or_cuda_str = "cpu" if cpu_or_cuda_str == "mps" else cpu_or_cuda_str
         assert (
             self.gradient_accumulation_steps > 0
         ), f"gradient_accumulation_steps must be > 0, not {self.gradient_accumulation_steps}"