From a27ab3de8e38d95c6a7400fd2032c6e352009ea0 Mon Sep 17 00:00:00 2001
From: Srini Iyer <sviyer@meta.com>
Date: Thu, 6 Feb 2025 00:07:59 +0000
Subject: [PATCH] Fix wandb logging

---
 bytelatent/metrics.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/bytelatent/metrics.py b/bytelatent/metrics.py
index 77dc4d7..e746e4f 100644
--- a/bytelatent/metrics.py
+++ b/bytelatent/metrics.py
@@ -4,7 +4,6 @@
 import json
 import logging
 from collections import namedtuple
-from dataclasses import asdict
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Union
@@ -68,8 +67,8 @@ class MetricLogger:
             and get_is_master()
         ):
             run = wandb.init(
-                config=asdict(self.args),
-                **asdict(self.args.logging.wandb),
+                config=self.args.model_dump(),
+                **self.args.logging.wandb.model_dump(),
             )
 
     def log(self, metrics: dict[str, Any]):