add requirements and fix minor bug in gpt2 example

logix-project · Feb 2, 2024 · fce30df · fce30df
1 parent c01fa48
commit fce30df
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/examples/gpt2_influence/train.py b/examples/gpt2_influence/train.py
@@ -22,7 +22,7 @@
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 train_loader, _, valid_loader = get_loaders()
-model = construct_model().to(device)
+model = construct_model()[0].to(device)
 
 
 def train(
@@ -100,7 +100,7 @@ def model_evaluate(model: nn.Module, loader: torch.utils.data.DataLoader) -> flo
     start_time = time.time()
 
     set_seed(i)
-    model = construct_model()
+    model, _ = construct_model()
 
     train(
         model=model,

diff --git a/requirements.txt b/requirements.txt
@@ -4,4 +4,7 @@ torch
 einops
 
 pyyaml
+transformers
+datasets
+accelerate