fix: adding wandb under 'full' section

Author: init-22

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
   "clu==0.0.12",
   "matplotlib>=3.9.2",
   "tabulate==0.9.0",
+
 ]
 
 [build-system]
@@ -70,7 +71,7 @@ version_file = "algorithmic_efficiency/_version.py"
 [project.optional-dependencies]
 # All workloads
 full = [
-  "algorithmic_efficiency[criteo1tb,fastmri,ogbg,librispeech_conformer,wmt]",
+  "algorithmic_efficiency[criteo1tb,fastmri,ogbg,librispeech_conformer,wmt,wandb]",
 ]
 # All workloads plus development dependencies
 full_dev = ["algorithmic_efficiency[full,dev]"]
@@ -83,6 +84,8 @@ dev = [
   "pre-commit==4.0.1",
 ]
 
+wandb = ["wandb==0.16.5"]
+
 # Workloads
 criteo1tb = ["scikit-learn==1.5.2"]
 fastmri = ["h5py==3.12.0", "scikit-image==0.24.0"]
@@ -119,7 +122,6 @@ pytorch_gpu = [
   "torch==2.5.1",
   "torchvision==0.20.1",
 ] # Note: omit the cuda suffix and installing from the appropriate wheel will result in using locally installed CUDA.
-wandb = ["wandb==0.16.5"]
 
 ###############################################################################
 #                           Linting Configurations                            #

submission_runner.py

@@ -21,7 +21,6 @@
 import itertools
 import json
 import os
-
 import struct
 import time
 from types import MappingProxyType
@@ -685,7 +684,7 @@ def main(_):
   base_workload = workloads.get_base_workload_name(FLAGS.workload)
   if base_workload == 'librispeech_conformer':
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
-  
+
   if FLAGS.set_pytorch_max_split_size:
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:256'