Corrected the example of value_and_grad (#274)

* Corrected the example for mx.value_and_grad

* Reformat through pre-commit/black
This commit is contained in:
Vidit Agarwal 2023-12-24 00:36:38 +05:30 committed by GitHub
parent f91f450141
commit acf1721b98
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 4 deletions

View File

@ -88,5 +88,3 @@ class Dropout2d(Module):
mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape)
return (1 / self._p_1) * mask * x

View File

@ -569,7 +569,7 @@ void init_transforms(py::module_& m) {
return lvalue
# Returns lvalue, dlvalue/dparams
lvalue, grads = mx.value_and_grad(mse)
lvalue, grads = mx.value_and_grad(mse)(params, inputs, targets)
def lasso(params, inputs, targets, a=1.0, b=1.0):
outputs = forward(params, inputs)
@ -580,7 +580,7 @@ void init_transforms(py::module_& m) {
return loss, mse, l1
(loss, mse, l1), grads = mx.value_and_grad(lasso)
(loss, mse, l1), grads = mx.value_and_grad(lasso)(params, inputs, targets)
Args:
fun (function): A function which takes a variable number of