| 276 | mx.grad(loss_fn)(model) |
| 277 | |
| 278 | def test_update(self): |
| 279 | m = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3)) |
| 280 | |
| 281 | # Updating non-existent parameters |
| 282 | with self.assertRaises(ValueError): |
| 283 | updates = {"layers": [{"value": 0}]} |
| 284 | m.update(updates) |
| 285 | |
| 286 | with self.assertRaises(ValueError): |
| 287 | updates = {"layers": ["hello"]} |
| 288 | m.update(updates) |
| 289 | |
| 290 | # Wronge type |
| 291 | with self.assertRaises(ValueError): |
| 292 | updates = {"layers": [{"weight": "hi"}]} |
| 293 | m.update(updates) |
| 294 | |
| 295 | def test_update_modules(self): |
| 296 | m = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3)) |