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