update compile compatible

test.py

@@ -20,6 +20,7 @@
 from tqdm import tqdm
 from model import ft_net, ft_net_dense, ft_net_hr, ft_net_swin, ft_net_swinv2, ft_net_efficient, ft_net_NAS, ft_net_convnext, PCB, PCB_test
 from utils import fuse_all_conv_bn
+version =  torch.__version__
 #fp16
 try:
     from apex.fp16_utils import *
@@ -280,7 +281,7 @@ def get_id(img_path):
 #if opt.fp16:
 #    model_structure = network_to_half(model_structure)
 
-if torch.cuda.get_device_capability()[0]>6 and len(opt.gpu_ids)==1: # should be >=7
+if torch.cuda.get_device_capability()[0]>6 and len(opt.gpu_ids)==1 and int(version[0])>1: # should be >=7
     print("Compiling model...")
     # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0
     torch.set_float32_matmul_precision('high')

train.py

@@ -588,7 +588,7 @@ def draw_curve(current_epoch):
     #optimizer_ft = FP16_Optimizer(optimizer_ft, static_loss_scale = 128.0)
     model, optimizer_ft = amp.initialize(model, optimizer_ft, opt_level = "O1")
 
-if torch.cuda.get_device_capability()[0]>6 and len(opt.gpu_ids)==1: # should be >=7 and one gpu
+if torch.cuda.get_device_capability()[0]>6 and len(opt.gpu_ids)==1 and int(version[0])>1: # should be >=7 and one gpu
     torch.set_float32_matmul_precision('high')
     print("Compiling model... The first epoch may be slow, which is expected!")
     # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0

utils.py

@@ -1,5 +1,6 @@
-import torch.nn as nn
 import os
+import torch
+import torch.nn as nn
 from torch.nn.utils import fuse_conv_bn_eval
 
 class CrossEntropyLabelSmooth(nn.Module):