From e23bca2493d5ff21efc54ce0eec585c5c107079e Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Thu, 25 Jul 2024 16:23:34 -0400 Subject: [PATCH] tf_keras in test_utils --- tests/test_modelscan.py | 2 +- tests/test_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 5fabedf..1c10a76 100644 --- 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 index 52a195c..15236d7 100644 --- 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