Fixed the clipping operation in fast gradient clipping with Privacy Engine (#664)

Summary:
Pull Request resolved: #664

Earlier, privacy engine erroneously wasn't using the max_grad_norm argument which infoms the clipping norm. It is fixed now.

Reviewed By: HuanyuZhang

Differential Revision: D60917356

fbshipit-source-id: aad3bcf08048027d93ae44400caf8d7601d90961

Author: EnayatUllah

opacus/privacy_engine.py

@@ -171,6 +171,7 @@ def _prepare_model(
         module: nn.Module,
         *,
         batch_first: bool = True,
+        max_grad_norm: Union[float, List[float]] = 1.0,
         loss_reduction: str = "mean",
         grad_sample_mode: str = "hooks",
     ) -> AbstractGradSampleModule:
@@ -194,12 +195,21 @@ def _prepare_model(
 
             return module
         else:
-            return wrap_model(
-                module,
-                grad_sample_mode=grad_sample_mode,
-                batch_first=batch_first,
-                loss_reduction=loss_reduction,
-            )
+            if grad_sample_mode == "ghost":
+                return wrap_model(
+                    module,
+                    grad_sample_mode=grad_sample_mode,
+                    batch_first=batch_first,
+                    loss_reduction=loss_reduction,
+                    max_grad_norm=max_grad_norm,
+                )
+            else:
+                return wrap_model(
+                    module,
+                    grad_sample_mode=grad_sample_mode,
+                    batch_first=batch_first,
+                    loss_reduction=loss_reduction,
+                )
 
     def is_compatible(
         self,
@@ -355,6 +365,7 @@ def make_private(
         module = self._prepare_model(
             module,
             batch_first=batch_first,
+            max_grad_norm=max_grad_norm,
             loss_reduction=loss_reduction,
             grad_sample_mode=grad_sample_mode,
         )