Update trace.py

trace.py

@@ -9,6 +9,7 @@ def trace_wandb(config,
                 prompt,
                 completion,
                 result,
+                callback,
                 err_msg,
                 start_time_ms,
                 end_time_ms):
@@ -24,6 +25,7 @@ def trace_wandb(config,
                   "config": str(config)
                  } if (str(err_msg) == "") else {},
         outputs = {"result": str(result),
+                   "callback": str(callback),
                    "completion": str(completion)
                   } if (str(err_msg) == "") else {},
         start_time_ms = start_time_ms,