From 7d9ac9a6439b43a0e6b1315a73e2332144125671 Mon Sep 17 00:00:00 2001
From: Siddhi Kiran Bajracharyya <siddhi.47.skb@gmail.com>
Date: Mon, 3 Apr 2023 18:00:58 -0500
Subject: [PATCH] Update custom_training_loop.py

formatted imports
---
 core/custom_training_loop.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/core/custom_training_loop.py b/core/custom_training_loop.py
index 0241145..57aba03 100644
--- a/core/custom_training_loop.py
+++ b/core/custom_training_loop.py
@@ -8,14 +8,22 @@ Features (go/100_enablements):
     - go/dataset-service 0-copy integration
 
 """
-import datetime
 import os
-from typing import Callable, Dict, Iterable, List, Mapping, Optional
+import datetime
+from typing import (
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional
+)
+
 
 
 from tml.common import log_weights
-import tml.common.checkpointing.snapshot as snapshot_lib
 from tml.core.losses import get_global_loss_detached
+import tml.common.checkpointing.snapshot as snapshot_lib
 from tml.ml_logging.torch_logging import logging  # type: ignore[attr-defined]
 from tml.core.train_pipeline import TrainPipelineSparseDist