feibyte
diff --git a/‎Dockerfile‎
Lines changed: 1 addition & 0 deletions b/‎Dockerfile‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎tests/test_transformers.py‎
Lines changed: 25 additions & 0 deletions b/‎tests/test_transformers.py‎
Lines changed: 25 additions & 0 deletions
@@ -481,6 +481,7 @@ RUN pip install flashtext && \
     pip install qgrid && \
     pip install bqplot && \
     pip install earthengine-api && \
+    pip install transformers && \
     /tmp/clean-layer.sh
 
 # Tesseract and some associated utility packages
 
@@ -0,0 +1,25 @@
+import unittest
+
+import torch
+from transformers import AdamW
+
+
+class TestTransformers(unittest.TestCase):
+    def assertListAlmostEqual(self, list1, list2, tol):
+        self.assertEqual(len(list1), len(list2))
+        for a, b in zip(list1, list2):
+            self.assertAlmostEqual(a, b, delta=tol)
+
+    def test_adam_w(self):
+        w = torch.tensor([0.1, -0.2, -0.1], requires_grad=True)
+        target = torch.tensor([0.4, 0.2, -0.5])
+        criterion = torch.nn.MSELoss()
+        # No warmup, constant schedule, no gradient clipping
+        optimizer = AdamW(params=[w], lr=2e-1, weight_decay=0.0)
+        for _ in range(100):
+            loss = criterion(w, target)
+            loss.backward()
+            optimizer.step()
+            w.grad.detach_() # No zero_grad() function on simple tensors. we do it ourselves.
+            w.grad.zero_()
+        self.assertListAlmostEqual(w.tolist(), [0.4, 0.2, -0.5], tol=1e-2)