tf_keras in test_utils

protectai · Jul 25, 2024 · e23bca2 · e23bca2
1 parent 992c466
commit e23bca2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py
@@ -16,7 +16,7 @@
 import tensorflow as tf
 import tf_keras as keras
 from typing import Any, List, Set, Dict
-from tests.test_utils import (
+from test_utils import (
     generate_dill_unsafe_file,
     generate_unsafe_pickle_file,
     MaliciousModule,

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -6,7 +6,7 @@
 import torch
 import torch.nn as nn
 import tensorflow as tf
-from tensorflow import keras
+import tf_keras as keras
 
 
 class MaliciousModule(keras.Model):  # type: ignore