From 75b1d8aaffccb53df53a7088113ce42b9d3e0906 Mon Sep 17 00:00:00 2001
From: thandal <than@timbrel.org>
Date: Sun, 22 Feb 2026 09:16:19 -0500
Subject: [PATCH] Fix wandb handling of single channel images in logger

Handle single channel images by replicating the channel.
---
 elements/logger.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/elements/logger.py b/elements/logger.py
index 245f5ba..5036d96 100644
--- a/elements/logger.py
+++ b/elements/logger.py
@@ -319,6 +319,8 @@ def __call__(self, summaries):
         bystep[step][name] = wandb.Image(value)
       elif len(value.shape) == 4:
         assert value.shape[3] in [1, 3, 4], value.shape
+        if value.shape[3] == 1:  # Address "ValueError: Can't write images with one color channel." issue
+          value = np.concatenate((value,) * 3, axis=3)
         value = np.transpose(value, [0, 3, 1, 2])
         if value.dtype != np.uint8:
           value = (255 * np.clip(value, 0, 1)).astype(np.uint8)