Browse Source

!3487 Modify comment of register_backward_hook

Merge pull request !3487 from Simson/push-to-opensource
tags/v0.7.0-beta
mindspore-ci-bot Gitee 5 years ago
parent
commit
57c00a77f1
2 changed files with 3 additions and 2 deletions
  1. +1
    -1
      mindspore/nn/cell.py
  2. +2
    -1
      mindspore/ops/operations/debug_ops.py

+ 1
- 1
mindspore/nn/cell.py View File

@@ -862,7 +862,7 @@ class Cell:

def register_backward_hook(self, fn):
"""
Set the cell backward hook function.
Set the cell backward hook function. Note that this function is only supported in Pynative Mode.

Note:
fn should be defined as following code shows, `cell_name` is the name of registered cell,


+ 2
- 1
mindspore/ops/operations/debug_ops.py View File

@@ -251,7 +251,8 @@ class InsertGradientOf(PrimitiveWithInfer):

class HookBackward(PrimitiveWithInfer):
"""
Used as tag to hook gradient in intermediate variables.
Used as tag to hook gradient in intermediate variables. Note that this function
is only supported in Pynative Mode.

Note:
The hook function should be defined like `hook_fn(grad) -> Tensor or None`,


Loading…
Cancel
Save