mirror of
https://github.com/baker-laboratory/RoseTTAFold-All-Atom.git
synced 2024-11-24 22:37:20 +00:00
6 lines
174 B
Python
6 lines
174 B
Python
|
# for gradient checkpointing
|
||
|
|
def create_custom_forward(module, **kwargs):
|
||
|
|
def custom_forward(*inputs):
|
||
|
|
return module(*inputs, **kwargs)
|
||
|
|
return custom_forward
|