refactor(mgb/grad): place grad at param.grad

GitOrigin-RevId: fddeace40248f2edf7781458ef408529d75022ca

imperative/python/megengine/autodiff/grad_manager.py

@@ -42,14 +42,15 @@ class GradManager:
                 self._recording = True
                 self._grad = grad
                 for params, callbacks in self._call_back_pair:
+                    for p in params:
 
-                    def callback(param, grad, callbacks=callbacks):
-                        ret = grad
-                        for cb in callbacks:
-                            ret = cb(param, ret)
-                        param.grad = ret
+                        def callback(param, grad, callbacks=callbacks, p=p):
+                            ret = grad
+                            for cb in callbacks:
+                                ret = cb(param, ret)
+                            p.grad = ret
 
-                    grad.wrt(*params, callback=callback)
+                        grad.wrt(p, callback=callback)
                 with grad:
                     yield
             finally:

imperative/python/megengine/optimizer/sgd.py

@@ -52,12 +52,6 @@ class SGD(Optimizer):
         momentum = param_group["momentum"]
 
         for param in param_group["params"]:
-
-            if not isinstance(param.grad, Buffer):
-                raise TypeError(
-                    "grad must be a Buffer, maybe you forget to call backward()?"
-                )
-
             if not param.requires_grad:
                 continue