small fixes

Co-authored-by: Copilot <copilot@github.com>

Author: zulissimeta

src/fairchem/core/units/mlip_unit/batch_server.py

@@ -185,6 +185,14 @@ def _split_predictions(
                         f"and num_atoms {batch.num_atoms}"
                     )
 
+                # Move to CPU before returning so the caller (which may be a
+                # CPU-only Ray worker) can deserialize the result without
+                # requiring CUDA.
+                if hasattr(system_predictions[key], "detach"):
+                    system_predictions[key] = (
+                        system_predictions[key].detach().cpu()
+                    )
+
             split_preds.append(system_predictions)
 
         return split_preds
@@ -226,6 +234,9 @@ def __init__(
             "BatchPredictServer initialized with predict_unit from object store"
         )
 
+    def is_multiplexed(self) -> bool:
+        return False
+
 
 @serve.deployment(
     logging_config=serve.schema.LoggingConfig(log_level="WARNING"),
@@ -270,6 +281,9 @@ def __init__(
         self.configure_batching(max_batch_size, batch_wait_timeout_s)
         logging.info("MultiplexedBatchPredictServer initialized")
 
+    def is_multiplexed(self) -> bool:
+        return True
+
     @serve.batch(
         batch_size_fn=lambda batch: sum(sample.natoms.sum() for sample in batch).item()
     )