From 6beddf9f9ea0f10deb0d3978593b7e0c56b23622 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Fri, 13 Feb 2026 06:51:55 +0000
Subject: [PATCH] fix: change epoch type from int to float to match
 TrainerState

---
 studio/backend/core/training/trainer.py | 2 +-
 studio/backend/models/training.py       | 2 +-
 studio/backend/routes/training.py       | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index af5c3ebe01..1466a28674 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -34,7 +34,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class TrainingProgress:
     """Training progress tracking"""
-    epoch: int = 0
+    epoch: float = 0
     step: int = 0
     total_steps: int = 0
     loss: float = 0.0
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 3afcb187e7..700b7569b4 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -99,7 +99,7 @@ class TrainingProgress(BaseModel):
     loss: float = Field(..., description="Current loss value")
     learning_rate: float = Field(..., description="Current learning rate")
     progress_percent: float = Field(..., description="Progress percentage (0.0 to 100.0)")
-    epoch: Optional[int] = Field(None, description="Current epoch")
+    epoch: Optional[float] = Field(None, description="Current epoch")
     elapsed_seconds: Optional[float] = Field(None, description="Time elapsed since training started")
     eta_seconds: Optional[float] = Field(None, description="Estimated time remaining")
     grad_norm: Optional[float] = Field(None, description="L2 norm of gradients, computed before gradient clipping")
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index ee36a9b450..eaf004f7b2 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -436,7 +436,7 @@ async def stream_training_progress(
             loss: float,
             learning_rate: float,
             total_steps: int,
-            epoch: Optional[int] = None,
+            epoch: Optional[float] = None,
         ) -> TrainingProgress:
             total = max(total_steps, 0)
             if step < 0 or total == 0: