From 79bdfac28afb526430a938d38513c46936f8670a Mon Sep 17 00:00:00 2001
From: Karthik Prasad <krp@meta.com>
Date: Tue, 1 Aug 2023 12:37:56 -0700
Subject: [PATCH] Type check lint fixes (#599)

Summary:
## Types of changes

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Docs change / refactoring / dependency upgrade

## Motivation and Context / Related issue

Fixing flake8 lint issue detected in https://app.circleci.com/pipelines/github/pytorch/opacus/2835/workflows/1318787a-42c8-42a2-aba5-85e32eac7d2a/jobs/16079

## How Has This Been Tested (if it applies)

CI

## Checklist

- [x] The documentation is up-to-date with the changes I made.
- [x] I have read the **CONTRIBUTING** document and completed the CLA (see **CONTRIBUTING**).
- [x] All tests passed, and additional code has been covered with new tests.

Pull Request resolved: https://github.com/pytorch/opacus/pull/599

Reviewed By: JohnlNguyen

Differential Revision: D47936830

Pulled By: karthikprasad

fbshipit-source-id: 800f18ebfaf10ff65e9ceb2a0dbb9f182efee759
---
 opacus/grad_sample/conv.py | 6 +++---
 opacus/privacy_engine.py   | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/opacus/grad_sample/conv.py b/opacus/grad_sample/conv.py
index 16f2d95b..633782f9 100644
--- a/opacus/grad_sample/conv.py
+++ b/opacus/grad_sample/conv.py
@@ -51,7 +51,7 @@ def compute_conv_grad_sample(
         return ret
 
     # get activations and backprops in shape depending on the Conv layer
-    if type(layer) == nn.Conv2d:
+    if type(layer) is nn.Conv2d:
         activations = unfold2d(
             activations,
             kernel_size=layer.kernel_size,
@@ -59,7 +59,7 @@ def compute_conv_grad_sample(
             stride=layer.stride,
             dilation=layer.dilation,
         )
-    elif type(layer) == nn.Conv1d:
+    elif type(layer) is nn.Conv1d:
         activations = activations.unsqueeze(-2)  # add the H dimension
         # set arguments to tuples with appropriate second element
         if layer.padding == "same":
@@ -77,7 +77,7 @@ def compute_conv_grad_sample(
             stride=(1, layer.stride[0]),
             dilation=(1, layer.dilation[0]),
         )
-    elif type(layer) == nn.Conv3d:
+    elif type(layer) is nn.Conv3d:
         activations = unfold3d(
             activations,
             kernel_size=layer.kernel_size,
diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py
index e9f59eb3..0ca6811b 100644
--- a/opacus/privacy_engine.py
+++ b/opacus/privacy_engine.py
@@ -181,7 +181,7 @@ def _prepare_model(
             if (
                 module.batch_first != batch_first
                 or module.loss_reduction != loss_reduction
-                or type(module) != get_gsm_class(grad_sample_mode)
+                or type(module) is not get_gsm_class(grad_sample_mode)
             ):
                 raise ValueError(
                     f"Pre-existing GradSampleModule doesn't match new arguments."