修改测试代码

train.py

@@ -10,9 +10,10 @@ import torchvision.transforms as transforms
 from matplotlib import pyplot as plt
 from torch import optim
 from tqdm import tqdm  # 导入tqdm
+
+import secret_func
 from model.Alexnet import Alexnet
 from watermark_codec import ModelEncoder
-from watermark_codec.tool import secret_func
 
 # 参数
 batch_size = 500

val.py

@@ -9,7 +9,7 @@ import torchvision.transforms as transforms
 
 from model.Alexnet import Alexnet
 from watermark_codec import ModelDecoder
-from watermark_codec.tool import secret_func
+import secret_func
 
 model_path = './run/train/alex_net.pt'
 key_path = './run/train/key.pt'