Figure out where the NaNs in autoencoder training are coming from:

[robogast]

Source:
2021-11-01/12-14-29/0/lightning_logs/version_8317429

[robogast]

TL;DR: instability probably happens in the VQ-layer, but still unsure what exactly happens.
Increasing commitment loss, and making sure cdist compute_mode is non-mm seems to at least mitigate the issue.
Forcing 32-bit with torch.cuda.amp.autocast(enabled=False) doesn't solve the issue.