[experimenting] use contiguous instead of realize in optim (#3770)

* run CI

* comment

* remove t.grad to try

* Revert "remove t.grad to try"

This reverts commit 05ec2d3b89.
This commit is contained in:
David Hou
2024-03-15 23:06:50 -07:00
committed by GitHub
parent e3e89c244b
commit 07324b56d5

View File

@@ -34,10 +34,9 @@ class SGD(Optimizer):
def step(self) -> None:
for i, t in enumerate(self.params):
assert t.grad is not None
# this is needed since the grads can form a "diamond"
# contiguous is needed since the grads can allegedly form a "diamond"
# TODO: fix this in lazy.py
t.grad.realize()
g = t.grad + self.wd * t.detach()
g = t.grad.contiguous() + self.wd * t.detach()
if self.momentum:
self.b[i].assign(self.momentum * self.b[i] + g) # NOTE: self.b[i] is zero on the first run, no if required
g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]