diff --git a/main.py b/main.py
index ff5b916..c13d659 100644
--- a/main.py
+++ b/main.py
@@ -1,8 +1,5 @@
-import os
 import wandb
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
 from argparse import ArgumentParser
 from time import time