zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 01:41:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix module attribute set, reset, set (#1403)

Browse Source
This commit is contained in:
Awni Hannun 2024-09-11 16:30:42 -07:00 committed by GitHub
parent 02efb310ca
commit 8b30acd7eb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
python/mlx/nn/layers/base.py
View File

@ -112,6 +112,7 @@ class Module(dict):
self[key] = val
else:
super(Module, self).__setattr__(key, val)
self.pop(key, None)
def load_weights(
self,

6
python/tests/test_nn.py
View File

@ -67,6 +67,12 @@ class TestBase(mlx_tests.MLXTestCase):
model = Model()
self.assertTrue(mx.array_equal(model.val, mx.array(1.0)))
model.val = None
self.assertEqual(model.val, None)
model.val = mx.array([3])
self.assertEqual(model.val.item(), 3)
def test_model_with_dict(self):
class DictModule(nn.Module):
def __init__(self):