fix

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Oct 30, 2024 · 227a693 · 227a693
1 parent cca1958
commit 227a693
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/collections/llm/gpt/model/starcoder.py b/nemo/collections/llm/gpt/model/starcoder.py
@@ -18,6 +18,7 @@
 
 import torch.nn.functional as F
 from torch import nn
+import torch
 
 from nemo.collections.llm.gpt.model.base import GPTConfig, GPTModel, torch_dtype_from_mcore_config
 from nemo.collections.llm.utils import Config