Fix DistributedDP Optimizer for Fast Gradient Clipping (#662)

Summary:
Pull Request resolved: #662

The step function incorrectly called "original_optimizer.original_optimizer" instead of  "original_optimizer". Fixed it now.

Reviewed By: HuanyuZhang

Differential Revision: D60484128

fbshipit-source-id: 1bde00292b2afccc31803ebefb2c361dc7e9bb77

Author: EnayatUllah

opacus/__init__.py

@@ -14,14 +14,15 @@
 # limitations under the License.
 
 from . import utils
-from .grad_sample import GradSampleModule
+from .grad_sample import GradSampleModule, GradSampleModuleFastGradientClipping
 from .privacy_engine import PrivacyEngine
 from .version import __version__
 
 
 __all__ = [
     "PrivacyEngine",
     "GradSampleModule",
+    "GradSampleModuleFastGradientClipping",
     "utils",
     "__version__",
 ]

opacus/optimizers/ddpoptimizer_fast_gradient_clipping.py

@@ -76,6 +76,6 @@ def step(
 
         if self.pre_step():
             self.reduce_gradients()
-            return self.original_optimizer.original_optimizer.step()
+            return self.original_optimizer.step()
         else:
             return None