diff --git a/main.py b/main.py
index 9c020fc..0c3ae52 100644
--- a/main.py
+++ b/main.py
@@ -28,7 +28,7 @@ def setupLogger():
     )
     logger.addHandler(logger_handler)
 
-def interrupt_handler(signum, driver: webdriver.Chrome):
+def interrupt_handler(signum, driver: webdriver.Remote):
     logger.info(f'Handling signal {signum} ({signal.Signals(signum).name}).')
     
     driver.quit()
@@ -157,7 +157,7 @@ def downloadStats(driver: webdriver.Chrome, peersDict: dict):
 
     saveStats([stats])
     
-def setupStats(driver: webdriver.Chrome, url: str):
+def setupStats(driver: webdriver.Remote, url: str):
     logger.log(logging.INFO, 'Setting up stats.')
     actions = ActionChains(driver)
     wait = WebDriverWait(driver, 30, poll_frequency=0.2)