Skip to content

Commit

Permalink
Update main.py
Browse files Browse the repository at this point in the history
hazdzz authored May 10, 2024
1 parent 6156441 commit c5293c8
Showing 1 changed file with 4 additions and 5 deletions.
9 changes: 4 additions & 5 deletions main.py
Original file line number Diff line number Diff line change
@@ -20,20 +20,19 @@

def set_env(seed):
# Set available CUDA devices
# This option is crucial for an multi-GPU device
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
# os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
os.environ['PYTHONHASHSEED']=str(seed)
# os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
os.environ['PYTHONHASHSEED'] = str(seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True)
# torch.use_deterministic_algorithms(True)

def get_parameters():
parser = argparse.ArgumentParser(description='STGCN')

0 comments on commit c5293c8

Please sign in to comment.