I’m working on a project forked from a HuggingFace code base that used NVIDIA’s apex.amp.initialize.
            try:
                from apex import amp
            except ImportError:
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
            self.student, self.optimizer = amp.initialize(
                self.student,
                self.optimizer,
                opt_level=self.run_config['devices']['fp16_opt_level']
            )

What’s the equivalent function call using torch.cuda.amp?