diff --git a/action.yml b/action.yml
index a9aeae2..5aae880 100644
--- a/action.yml
+++ b/action.yml
@@ -22,7 +22,7 @@ runs:
       with:
         python-version: 3.x
     - name: Clone moodle uploader
-      run: git clone https://github.com/RokasPuzonas/ktu-moodle-assignment-upload -b v1.0.1
+      run: git clone https://github.com/RokasPuzonas/ktu-moodle-assignment-upload -b v1.0.2
       shell: "bash"
     - name: Setup Geckodriver
       uses: browser-actions/setup-geckodriver@v0.0.0
diff --git a/main.py b/main.py
index 75bc1fa..08c1f89 100755
--- a/main.py
+++ b/main.py
@@ -2,6 +2,8 @@
 from selenium import webdriver
 import click
 from selenium.webdriver.firefox.options import Options
+from selenium.webdriver.remote.webdriver import WebDriver
+from selenium.common.exceptions import NoSuchElementException
 from selenium.webdriver.common.by import By
 from dotenv import load_dotenv
 import time
@@ -12,12 +14,12 @@ load_dotenv()
 LOGIN_URL = "https://moodle.ktu.edu/login/index.php"
 EDIT_ASSIGNMENT_URL = "https://moodle.ktu.edu/mod/assign/view.php?id={id}&action=editsubmission"
 
-def create_driver():
+def create_driver() -> WebDriver:
     options = Options()
     options.headless = True
     return webdriver.Firefox(options=options)
 
-def login(driver, username, password):
+def login(driver: WebDriver, username, password):
     driver.get(LOGIN_URL)
 
     driver.find_element(By.ID, "username").send_keys(username)
@@ -28,15 +30,23 @@ def login(driver, username, password):
     driver.find_element(By.ID, "yesbutton").click()
     time.sleep(1)
 
-def upload_file_to_assignment(driver, assignment, filename, upload_filename):
+def safe_find_element(driver, *args, **kvargs):
+    try:
+        return driver.find_element(*args, **kvargs)
+    except NoSuchElementException:
+        return None
+
+def upload_file_to_assignment(driver: WebDriver, assignment, filename, upload_filename):
     if not upload_filename:
         upload_filename = path.basename(filename)
 
     driver.get(EDIT_ASSIGNMENT_URL.format(id=assignment))
     time.sleep(1)
 
-    file = driver.find_element(By.XPATH, f"//*[text()='{upload_filename}']")
+    # Check if file exists
+    file = safe_find_element(driver, By.XPATH, f"//*[text()='{upload_filename}']")
     if file:
+        # If it does, delete it
         file.click()
         time.sleep(1)
         driver.find_element(By.XPATH, "//button[text()='Naikinti']").click()