tf32推理

scripts/evaluation/world_model_interaction.py

@@ -18,6 +18,9 @@ from collections import OrderedDict
 from torch import nn
 from eval_utils import populate_queues, log_to_tensorboard
 from collections import deque
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
 from torch import Tensor
 from torch.utils.tensorboard import SummaryWriter
 from PIL import Image