Update train.py
Browse files
train.py
CHANGED
|
@@ -158,6 +158,7 @@ wandb.init(
|
|
| 158 |
# ============================
|
| 159 |
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
|
| 160 |
h_states = None
|
|
|
|
| 161 |
|
| 162 |
for epoch in range(num_epochs):
|
| 163 |
print(f"Starting Epoch {epoch + 1}/{num_epochs}")
|
|
|
|
| 158 |
# ============================
|
| 159 |
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
|
| 160 |
h_states = None
|
| 161 |
+
step = 0
|
| 162 |
|
| 163 |
for epoch in range(num_epochs):
|
| 164 |
print(f"Starting Epoch {epoch + 1}/{num_epochs}")
|