Можно ли выполнить режим графика в pytorch?

Question

Я только что понял, насколько сильна графовая мода тензорного потока. Можно реализовать режим графика в тензорном потоке через

@tf.function
def train_once(args):
    ....

def train():
    for epoch in range(100):
        train_once()

, и это быстрее, чем код без'@tf.function'.

Мой вопрос заключается в том, возможно ли это в pytorch. Как выполнить режим графика в pytorch?