diff --git a/examples/dft_driver.py b/examples/dft_driver.py
index cdb96d9a..ef7a0923 100644
--- a/examples/dft_driver.py
+++ b/examples/dft_driver.py
@@ -20,11 +20,12 @@
 from gpu4pyscf.dft import rks, uks
 
 parser = argparse.ArgumentParser(description='Run DFT with GPU4PySCF for molecules')
-parser.add_argument("--input",    type=str,  default='benzene/coord')
-parser.add_argument("--basis",    type=str,  default='def2-tzvpp')
-parser.add_argument("--auxbasis", type=str,  default='def2-tzvpp-jkfit')
-parser.add_argument("--xc",       type=str,  default='B3LYP')
-parser.add_argument("--solvent",  type=str, default='')
+parser.add_argument("--input",        type=str,  default='benzene/coord')
+parser.add_argument("--basis",        type=str,  default='def2-tzvpp')
+parser.add_argument("--auxbasis",     type=str,  default='def2-tzvpp-jkfit')
+parser.add_argument("--xc",           type=str,  default='B3LYP')
+parser.add_argument("--solvent",      type=str,  default='')
+parser.add_argument('--unrestricted', type=bool, default=False)
 args = parser.parse_args()
 
 lib.num_threads(16)
@@ -37,8 +38,10 @@
 # set verbose >= 6 for debugging timer
 mol.verbose = 7
 
-mf_df = uks.UKS(mol, xc=args.xc).density_fit(auxbasis=args.auxbasis)
-#mf_df = rks.RKS(mol, xc=args.xc).density_fit(auxbasis=args.auxbasis)
+if args.unrestricted:
+    mf_df = uks.UKS(mol, xc=args.xc).density_fit(auxbasis=args.auxbasis)
+else:
+    mf_df = rks.RKS(mol, xc=args.xc).density_fit(auxbasis=args.auxbasis)
 mf_df.verbose = 7
 
 if args.solvent: