diff --git a/gasket_core.c b/gasket_core.c
index 5ae3d44..94e64b9 100644
--- a/gasket_core.c
+++ b/gasket_core.c
@@ -2089,9 +2089,9 @@ struct device *gasket_get_device(struct gasket_dev *dev)
  **/
 int gasket_wait_with_reschedule(
 	struct gasket_dev *gasket_dev, int bar, u64 offset, u64 mask, u64 val,
-	u64 max_retries, u64 delay_ms)
+	uint max_retries, u64 delay_ms)
 {
-	u64 retries = 0;
+	uint retries = 0;
 	u64 tmp;
 
 	while (retries < max_retries) {
diff --git a/gasket_core.h b/gasket_core.h
index 94a5537..50ad0c8 100644
--- a/gasket_core.h
+++ b/gasket_core.h
@@ -702,6 +702,6 @@ struct device *gasket_get_device(struct gasket_dev *dev);
 /* Helper function, Asynchronous waits on a given set of bits. */
 int gasket_wait_with_reschedule(
 	struct gasket_dev *gasket_dev, int bar, u64 offset, u64 mask, u64 val,
-	u64 max_retries, u64 delay_ms);
+	uint max_retries, u64 delay_ms);
 
 #endif /* __GASKET_CORE_H__ */