training: more efficient inference

training/infer.py

@@ -157,7 +157,7 @@ def save_images(path, image, image_srgb, num_channels, feature_ext=infer.main_fe
       else:
         save_image(filename_prefix + format, image_srgb, num_channels=num_channels[feature_ext])
 
-  with torch.no_grad():
+  with torch.inference_mode():
     for group, input_names, target_name in image_sample_groups:
       # Create the output directory if it does not exist
       output_group_dir = os.path.join(output_dir, os.path.dirname(group))

training/preprocess.py

@@ -139,7 +139,7 @@ def preprocess_dataset(data_name):
     save_config(output_dir, cfg)
 
   # Preprocess all datasets
-  with torch.no_grad():
+  with torch.inference_mode():
     for dataset in [cfg.train_data, cfg.valid_data]:
       if dataset:
         preprocess_dataset(dataset)

training/train.py

@@ -244,7 +244,7 @@ def main_worker(rank, cfg):
       valid_loss = 0.
 
       # Iterate over the batches
-      with torch.no_grad():
+      with torch.inference_mode():
         for _, batch in enumerate(valid_loader, 0):
           if cfg.device == 'cuda' and cfg.compile in {'reduce-overhead', 'max-autotune'}:
             torch.compiler.cudagraph_mark_step_begin()