mypy fix

Signed-off-by: sahil suneja <[email protected]>
foundation-model-stack · Sep 9, 2024 · ac548c5 · ac548c5
1 parent e08628d
commit ac548c5
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/speculator/train_speculator.py b/speculator/train_speculator.py
@@ -2,7 +2,7 @@
 import os
 import time
 
-import fire
+import fire  # type: ignore
 import torch
 import torch.optim as optim
 from fms.models import get_model

diff --git a/speculator/train_speculator_utils.py b/speculator/train_speculator_utils.py
@@ -313,7 +313,6 @@ def train_speculator(
     model.eval()
     speculator.train()
     ddp_stats = torch.zeros(2 + speculator.n_predict).to(local_rank)
-    train_loss = 0
 
     start = time.time()
     loop_start = time.time()