From 2291d49b22390de22cfb027ac14d56ce41847adc Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 10 Mar 2023 11:53:44 +0100 Subject: [PATCH] Update. --- main.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/main.py b/main.py index c01cc8f..5b18985 100755 --- a/main.py +++ b/main.py @@ -5,6 +5,9 @@ # Written by Francois Fleuret +# torch.backends.cuda.matmul.allow_tf23 +# torch.autocast(torch.bfloat16) + import math, sys, argparse, time, tqdm, itertools, os import torch, torchvision @@ -15,7 +18,11 @@ import mygpt, tensorstack ###################################################################### -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +if torch.cuda.is_available(): + device = torch.device("cuda") + torch.backends.cuda.matmul.allow_tf32 = True +else: + device = torch.device("cpu") ###################################################################### @@ -55,8 +62,6 @@ parser.add_argument("--nb_blocks", type=int, default=12) parser.add_argument("--dropout", type=float, default=0.1) -parser.add_argument("--nb_oneshot_blocks", type=int, default=-1) - parser.add_argument("--deterministic_synthesis", action="store_true", default=False) parser.add_argument("--no_checkpoint", action="store_true", default=False) -- 2.39.5