Address Issue #227

Hello! 8BitAdamW has stopped working for Dreambooth since 0.36.0 (loss NaNs very quickly). I took a stab at reviewing the differences between the releases. The if-block around functional.py:218 in 2f2063b may have been accidentally indented. tests.test_functional.test_few_bit_quant doesn't cover the inside of the if block, so maybe this was missed (or it was intentional and it'd be great to have confirmation). I'm hoping one of the maintainers can take a quick look. Thanks!