minor fix 🐛

muditbhargava66 · Apr 26, 2024 · 72b5f08 · 72b5f08
1 parent 113d2cb
commit 72b5f08
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/examples/vit_experiments/train.py b/examples/vit_experiments/train.py
@@ -40,6 +40,9 @@ def train(model, optimizer, criterion, train_loader, test_loader, epochs, device
                 outputs = model(images)
                 loss = criterion(outputs, labels)
 
+            # Clip gradients to prevent explosion
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
             scaler.scale(loss).backward()
 
             if isinstance(optimizer, DropGrad):