thanks a lot ! I just saw that now and in the meantime I implemented this on my end this morning which seems to work as well
                        with autocast():
                            outputs = model(inputs)
                            loss = criterion(outputs.squeeze(), labels)
                        self._scaler.scale(loss).backward()
                        self._scaler.unscale_(optimizer)
                        optimizer.first_step(zero_grad=True)

                        # second forward-backward pass
                        with autocast():
                            outputs = model(inputs)
                            criterion(outputs.squeeze(), labels).backward()
                        
                        optimizer.second_step(zero_grad=True)
                        self._scaler.update()