diff --git a/scripts/watch_distill.py b/scripts/watch_distill.py
index ef891ef..580027d 100644
--- a/scripts/watch_distill.py
+++ b/scripts/watch_distill.py
@@ -45,6 +45,9 @@ def parse_loss(line):
     if m:
         return float(m.group(1))
     m = re.search(r"loss[=:]\s*([\d.]+)", line)
+    if m:
+        return float(m.group(1))
+    m = re.search(r"loss=([\d.]+)", line)
     if m:
         return float(m.group(1))
     return None
@@ -54,6 +57,9 @@ def parse_step(line):
     m = re.search(r"(\d+)/(\d+)\s+\[", line)
     if m:
         return int(m.group(1)), int(m.group(2))
+    m = re.search(r"step=(\d+)", line)
+    if m:
+        return int(m.group(1)), None
     return None, None
 
 def parse_progress(line):
@@ -173,4 +179,4 @@ def main():
         sys.exit(0)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()