Как при использовании torch.autocast заставить отдельные слои плавать32? ⇐ Python

[Anonymous]

Я пытаюсь обучить модель со смешанной точностью. Однако я хочу, чтобы некоторые слои имели полную точность из соображений стабильности. Как заставить отдельный слой иметь тип float32 при использовании torch.autocast? В частности, мне бы хотелось, чтобы это можно было компилировать на onnx.
Это что-то вроде:

Код: Выделить всё

with torch.autocast(device_type='cuda', enabled=False, dtype=torch.float16):
out = my_unstable_layer(inputs.float())

Изменить:
Похоже, это действительно официальный метод. См. документацию по факелу.

Подробнее здесь: https://stackoverflow.com/questions/734 ... to-float32