From 8d1119dff18506588604949d2cc81997d63cd911 Mon Sep 17 00:00:00 2001 From: Gustaf Ahdritz Date: Thu, 18 Nov 2021 00:10:53 -0500 Subject: [PATCH] Update TorchScript call in inference script --- run_pretrained_openfold.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/run_pretrained_openfold.py b/run_pretrained_openfold.py index e0631548..f5cb3bf4 100644 --- a/run_pretrained_openfold.py +++ b/run_pretrained_openfold.py @@ -31,13 +31,12 @@ from openfold.config import model_config from openfold.data import templates, feature_pipeline, data_pipeline from openfold.model.model import AlphaFold -from openfold.model.primitives import Attention, GlobalAttention +from openfold.model.torchscript import script_primitives_ from openfold.np import residue_constants, protein import openfold.np.relax.relax as relax from openfold.utils.import_weights import ( import_jax_weights_, ) -from openfold.utils.torchscript_utils import script_submodules_ from openfold.utils.tensor_utils import ( tensor_tree_map, ) @@ -45,10 +44,6 @@ from scripts.utils import add_data_args -def script_primitives_(model): - script_submodules_(model, [Attention, GlobalAttention]) - - def main(args): config = model_config(args.model_name) model = AlphaFold(config)