From 638bd3dc075eb44e37595f89b7cb737ec9c1ef92 Mon Sep 17 00:00:00 2001
From: Michael McKinsey <MichaelMcKinsey1@gmail.com>
Date: Thu, 29 Jan 2026 14:45:12 -0800
Subject: [PATCH] Update worker.py

---
 ScaFFold/worker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py
index 07bedf0..431c2eb 100644
--- a/ScaFFold/worker.py
+++ b/ScaFFold/worker.py
@@ -187,7 +187,7 @@ def main(kwargs_dict: dict = {}):
             shard_dim=shard_dim,
             device_type=device.type,
         )
-        model = model.to(device).to(memory_format=torch.contiguous_format)
+        model = model.to(device, memory_format=torch.channels_last_3d)
         # Wrap with DistConvDDP that corrects gradient scaling for dc submesh
         model = DistConvDDP(
             model,