File tree Expand file tree Collapse file tree 2 files changed +4
-2
lines changed Expand file tree Collapse file tree 2 files changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -220,7 +220,10 @@ def _callback_lookup_(op):
220
220
:return: callback function
221
221
"""
222
222
if op .type == 'parallel_do' and op .attr ('use_nccl' ):
223
+ all_vars = op .block .vars
223
224
param_names = set (op .input ('parameters' ))
225
+ param_names = filter (lambda name : all_vars [name ].stop_gradient is False ,
226
+ param_names )
224
227
param_grad_names = [n + "@GRAD" for n in param_names ]
225
228
226
229
class ParallelDoCallBack (object ):
Original file line number Diff line number Diff line change @@ -294,8 +294,7 @@ def get_parameters(self):
294
294
295
295
params = list (set (params ))
296
296
297
- param_list = [parent_block .var (name ) for name in params ]
298
- return filter (lambda param : param .stop_gradient is False , param_list )
297
+ return [parent_block .var (name ) for name in params ]
299
298
300
299
def complete_op (self ):
301
300
main_program = self .helper .main_program
You can’t perform that action at this time.
0 commit comments