Merge pull request #64 from ROCmSoftwarePlatform/IFU-master-2021-12-08

IFU-master-2021-12-08

apex/transformer/amp/__init__.py

+from apex.transformer.amp.grad_scaler import GradScaler
+
+
+__all__ = [
+    "GradScaler",
+]

apex/transformer/enums.py

 # coding=utf-8
-# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

apex/transformer/functional/__init__.py

-from .fused_softmax import FusedScaleMaskSoftmax
+from apex.transformer.functional.fused_softmax import FusedScaleMaskSoftmax
 
 __all__ = [
     "FusedScaleMaskSoftmax",

apex/transformer/pipeline_parallel/__init__.py

+from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func
+from apex.transformer.pipeline_parallel.schedules.common import build_model
+
+
+__all__ = [
+    "get_forward_backward_func",
+    "build_model",
+]