Pass ruff check

PrimeIntellect-ai · Jan 13, 2025 · 92f24b1 · 92f24b1
1 parent a90b58b
commit 92f24b1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/zeroband/data.py b/src/zeroband/data.py
@@ -7,7 +7,6 @@
 from zeroband.utils.logging import get_logger
 
 import torch
-from torch.utils.data import DataLoader
 from torch.utils.data import IterableDataset, Dataset
 from torchdata.stateful_dataloader import StatefulDataLoader
 from torch.distributed.checkpoint.stateful import Stateful

diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -4,7 +4,6 @@
 from multiprocessing.process import _children # type: ignore
 
 import torch
-from torch.nn import functional as F
 import torch.distributed as dist
 from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy # type: ignore
 from torch.autograd.profiler import record_function
@@ -509,7 +508,7 @@ def pretty_dict(d, indent=2):
             else:
                 logger.debug(" " * indent + f"{key}: {value}")
 
-    logger.debug(f"config:")
+    logger.debug("config:")
     pretty_dict(config.model_dump())
 
     try:
@@ -537,7 +536,7 @@ def pretty_dict(d, indent=2):
                 prof.__exit__(None, None, None)
 
             logger.info("Exporting chrome trace.")
-            prof.export_chrome_trace(f"logs/profile.json.gz")
+            prof.export_chrome_trace("logs/profile.json.gz")
 
             width = 30
             logger.info("\n" + "*" * width + " GPU TIME " + "*" * width)