Include loss scaling in README code example (#523)

f0k · mcarilli · commit e87b5799e697 · 2019-10-07T17:10:54.000-07:00
diff --git a/README.md b/README.md
@@ -67,6 +67,9 @@ model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
 
 # Train your model
 ...
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+    scaled_loss.backward()
+...
 
 # Save checkpoint
 checkpoint = {
