mirror of
https://github.com/baker-laboratory/RoseTTAFold-All-Atom.git
synced 2024-09-15 22:08:31 +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
|