X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=e741094c1543483514aae45e391d1d3573030a2b;hb=8f806c62be589d2837f88ffca084ed2ae833124c;hp=c01cc8f3dc3653d64fb465d83795fd17adad5936;hpb=eea23df18f107fc65c810261c7775a9393ef7c8e;p=picoclvr.git

diff --git a/main.py b/main.py
index c01cc8f..e741094 100755
--- a/main.py
+++ b/main.py
@@ -5,6 +5,9 @@
 
 # Written by Francois Fleuret <francois@fleuret.org>
 
+# 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)
@@ -89,7 +94,7 @@ except FileExistsError:
         print(f"result directory {args.result_dir} already exists")
         exit(1)
 
-log_file = open(os.path.join(args.result_dir, args.log_filename), "w")
+log_file = open(os.path.join(args.result_dir, args.log_filename), "a")
 
 if args.seed >= 0:
     # torch.backends.cudnn.deterministic = True