Fix plot (#782)

* Fix plot

* Updates for torch.GradScaler

Author: AlexeyKozhevin

batchflow/models/torch/base.py

@@ -683,7 +683,7 @@ def make_infrastructure(self):
         self.make_loss()
         self.make_optimizer()
         self.make_decay()
-        self.scaler = torch.cuda.amp.GradScaler()
+        self.scaler = torch.GradScaler("cuda")
 
         self.setup_gradient_clipping()
         self.setup_weights_averaging()

batchflow/plotter/plot.py

@@ -855,10 +855,10 @@ def clear(self):
 
         self.annotations = {}
 
-        self.ax.clear()
         for layer in self.layers:
             for obj in layer.objects:
                 obj.remove()
+        self.ax.clear()
         self.layers = []