Use TRRosetta to predict protein structure for a sequence

[heavenAsk]

Hello, How can we use TRRosetta to predict protein structure for a sequence? Just use it to inference not for training. The following code seems not work.

from tape import TRRosetta
from tape.datasets import TRRosettaDataset

# Download data and place it under `<data_path>/trrosetta`

train_data = TRRosettaDatset('<data_path>', 'train')  # will subsample MSAs
valid_data = TRRosettaDatset('<data_path>', 'valid')  # will not subsample MSAs

model = TRRosetta.from_pretrained('xaa')  # valid choices are 'xaa', 'xab', 'xac', 'xad', 'xae'. Each corresponds to one of the ensemble models.

batch = train_data.collate_fn([train_data[0]])
loss, predictions = model(**batch)