diff --git a/src/openzeppelin/upgrades/Proxy.cairo b/src/openzeppelin/upgrades/Proxy.cairo
index 062fe95aa..026e792f9 100644
--- a/src/openzeppelin/upgrades/Proxy.cairo
+++ b/src/openzeppelin/upgrades/Proxy.cairo
@@ -2,7 +2,6 @@
 # OpenZeppelin Contracts for Cairo v0.2.0 (upgrades/Proxy.cairo)
 
 %lang starknet
-#%builtins pedersen range_check bitwise
 
 from starkware.cairo.common.cairo_builtins import HashBuiltin
 from starkware.starknet.common.syscalls import (
diff --git a/src/openzeppelin/upgrades/library.cairo b/src/openzeppelin/upgrades/library.cairo
index d27e03680..ac4ccf03d 100644
--- a/src/openzeppelin/upgrades/library.cairo
+++ b/src/openzeppelin/upgrades/library.cairo
@@ -120,8 +120,10 @@ namespace Proxy:
             range_check_ptr
         }(new_implementation: felt):
         with_attr error_message("Proxy: implementation hash cannot be zero"):
-            Proxy_implementation_hash.write(new_implementation)
+            assert_not_zero(new_implementation)
         end
+
+        Proxy_implementation_hash.write(new_implementation)
         Upgraded.emit(new_implementation)
         return ()
     end
diff --git a/tests/upgrades/test_upgrades.py b/tests/upgrades/test_upgrades.py
index f9e2fc71a..1ab9b078b 100644
--- a/tests/upgrades/test_upgrades.py
+++ b/tests/upgrades/test_upgrades.py
@@ -203,6 +203,26 @@ async def test_upgrade_from_non_admin(proxy_factory):
     )
 
 
+@pytest.mark.asyncio
+async def test__set_implementation_as_zero(proxy_factory):
+    admin, _, proxy, _, _ = proxy_factory
+
+    # initialize implementation
+    await signer.send_transaction(
+        admin, proxy.contract_address, 'initializer', [
+            admin.contract_address
+        ]
+    )
+
+    # upgrade should revert
+    await assert_revert(
+        signer.send_transaction(
+            admin, proxy.contract_address, 'upgrade', [0]
+        ),
+        reverted_with="Proxy: implementation hash cannot be zero"
+    )
+
+
 @pytest.mark.asyncio
 async def test_implementation_v2(after_upgrade):
     admin, _, proxy, _, v2_decl = after_upgrade