From 7cac3d28c5ce521f254e2b31df9ff2b56849660f Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Sat, 4 Oct 2025 08:04:42 -0700
Subject: [PATCH] cleanup

---
 dreamer4/dreamer4.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 067fd66..290c409 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+from collections import namedtuple
 from functools import partial
 
 import torch
@@ -49,6 +50,8 @@ except ImportError:
 
 LinearNoBias = partial(Linear, bias = False)
 
+TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips'))
+
 # helpers
 
 def exists(v):
@@ -780,7 +783,7 @@ class VideoTokenizer(Module):
 
         losses = (recon_loss, lpips_loss)
 
-        return total_loss, losses
+        return total_loss, TokenizerLosses(losses)
 
 # dynamics model, axial space-time transformer