Update util.py
This commit is contained in:
parent
bb8f4f2dc1
commit
a22c61800d
|
@ -109,7 +109,7 @@ def checkpoint(func, inputs, params, flag):
|
||||||
explicitly take as arguments.
|
explicitly take as arguments.
|
||||||
:param flag: if False, disable gradient checkpointing.
|
:param flag: if False, disable gradient checkpointing.
|
||||||
"""
|
"""
|
||||||
if False: # disabled checkpointing to allow requires_grad = False for main model
|
if flag: # disabled checkpointing to allow requires_grad = False for main model
|
||||||
args = tuple(inputs) + tuple(params)
|
args = tuple(inputs) + tuple(params)
|
||||||
return CheckpointFunction.apply(func, len(inputs), *args)
|
return CheckpointFunction.apply(func, len(inputs), *args)
|
||||||
else:
|
else:
|
||||||
|
|
Loading…
Reference in New Issue