diff --git a/tensorflow_privacy/privacy/bolt_on/__init__.py b/tensorflow_privacy/privacy/bolt_on/__init__.py
index 2fed767..4271f08 100644
--- a/tensorflow_privacy/privacy/bolt_on/__init__.py
+++ b/tensorflow_privacy/privacy/bolt_on/__init__.py
@@ -12,14 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """BoltOn Method for privacy."""
+
+from distutils import version
 import sys
-from distutils.version import LooseVersion
+
 import tensorflow.compat.v1 as tf
 
-if LooseVersion(tf.__version__) < LooseVersion("2.0.0"):
+if version.LooseVersion(tf.__version__) < version.LooseVersion("2.0.0"):
   raise ImportError("Please upgrade your version "
                     "of tensorflow from: {0} to at least 2.0.0 to "
-                    "use privacy/bolt_on".format(LooseVersion(tf.__version__)))
+                    "use privacy/bolt_on".format(
+                        version.LooseVersion(tf.__version__)))
 if hasattr(sys, "skip_tf_privacy_import"):  # Useful for standalone scripts.
   pass
 else:
diff --git a/tensorflow_privacy/privacy/logistic_regression/datasets.py b/tensorflow_privacy/privacy/logistic_regression/datasets.py
index 2b54a26..adf53b2 100644
--- a/tensorflow_privacy/privacy/logistic_regression/datasets.py
+++ b/tensorflow_privacy/privacy/logistic_regression/datasets.py
@@ -28,8 +28,9 @@ Includes two types of datasets:
 - MNIST 10-class classification dataset.
 """
 
-from typing import Tuple
 import dataclasses
+from typing import Tuple
+
 import numpy as np
 from sklearn import preprocessing
 import tensorflow as tf
diff --git a/tensorflow_privacy/privacy/privacy_tests/secret_sharer/generate_secrets.py b/tensorflow_privacy/privacy/privacy_tests/secret_sharer/generate_secrets.py
index f33b718..7fdd961 100644
--- a/tensorflow_privacy/privacy/privacy_tests/secret_sharer/generate_secrets.py
+++ b/tensorflow_privacy/privacy/privacy_tests/secret_sharer/generate_secrets.py
@@ -13,10 +13,11 @@
 # limitations under the License.
 """Generate random sequences."""
 
+import dataclasses
 import itertools
 import string
 from typing import Dict, List
-from dataclasses import dataclass
+
 import numpy as np
 
 
@@ -66,7 +67,7 @@ def generate_random_sequences(vocab: List[str],
   return list(seq)
 
 
-@dataclass
+@dataclasses.dataclass
 class SecretConfig:
   """Configuration of secret for secrets sharer.
 
@@ -85,7 +86,7 @@ class SecretConfig:
   num_references: int
 
 
-@dataclass
+@dataclasses.dataclass
 class Secrets:
   """Secrets for secrets sharer.