From e3e933c6bc1b3057b2a17e001e4cf5fcb1c9902b Mon Sep 17 00:00:00 2001 From: Chunyang Wen Date: Thu, 11 Jan 2024 03:23:42 +0800 Subject: [PATCH] Add type hint for Module (#412) --- python/mlx/nn/layers/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index 646f5f2dc..094c89326 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -54,6 +54,8 @@ class Module(dict): mx.eval(model.parameters()) """ + __call__: Callable + def __init__(self): """Should be called by the subclasses of ``Module``.""" self._no_grad = set()