Fix plot (#782)

* Fix plot

* Updates for torch.GradScaler

Author: AlexeyKozhevin

README.md

@@ -3,11 +3,7 @@
 [![PyTorch](https://img.shields.io/badge/PyTorch-2.0-orange.svg)](https://pytorch.org)
 [![codecov](https://codecov.io/gh/analysiscenter/batchflow/branch/master/graph/badge.svg)](https://codecov.io/gh/analysiscenter/batchflow)
 [![PyPI](https://badge.fury.io/py/batchflow.svg)](https://badge.fury.io/py/batchflow)
-<<<<<<< HEAD
 [![Status](https://github.com/analysiscenter/batchflow/actions/workflows/status.yml/badge.svg?branch=master)](https://github.com/analysiscenter/batchflow/actions?query=workflow%3Astatus)
-=======
-[![Status](https://github.com/analysiscenter/batchflow/workflows/status.yml/badge.svg?branch=master)](https://github.com/analysiscenter/batchflow/actions?query=workflow%3Astatus)
->>>>>>> e81a745f (Update status badge)
 
 
 # BatchFlow

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 = []