[FEATURE-REQUEST] Switch to AMP precision

[AJDERS]

Switching to AMP precision promises huge speed-ups for both training and inference, and can be done easily with pytorch.

Just doing something like:

with torch.autocast():
   output = model(input)
   loss = loss_fn(output, target)

Can lead to speed-ups of 50%-150%.

[AJDERS]

Forgot to mention: At the cost of almost no performance drop.