Как переписать строку torch.argmax(model_out[0], dim=-1) без использования torch?

Question

[LakeForest]

Не хочется грузить библиотеку ради двух строк...

model = transformers.AutoModelForTokenClassification.from_pretrained('model').to(device)
...
 with torch.no_grad():
        model_out = model(**encoded)
 predicted_tokens = torch.argmax(model_out[0], dim=-1)

Comments

[WhiteApfel]

В pycharm можно перейти к исходнику функции через F4, и оттуда скопипастить принцип действия. Если не получится, могу глянуть ближе к ночи

Answer 1

[imageman]

Не хочется грузить библиотеку torch - это как? Вы же её используете по полной? Через что вы вычисляете значения модели? Ведь строка "with torch.no_grad()" неспроста у вас.

По какой причине решено экономить на спичках? Время работы, экономия памяти?

Попробуй вместо "проблемного" "torch.argmax(model_out[0], dim=-1)" поставить 0 и проанализируй как поменялось использование ресурсов.