diff --git a/build.gradle b/build.gradle
index fe2971cf..f16137e5 100644
--- a/build.gradle
+++ b/build.gradle
@@ -92,7 +92,7 @@ ext.deps = [
 
 allprojects {
   group = "com.linkedin.tony"
-  project.version = "0.5.1"
+  project.version = "0.5.2"
 }
 
 task sourcesJar(type: Jar) {
diff --git a/tony-core/src/main/java/com/linkedin/tony/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
index 3156549a..bd0cd9d1 100644
--- a/tony-core/src/main/java/com/linkedin/tony/TonySession.java
+++ b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
@@ -14,7 +14,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -67,7 +66,7 @@ public class TonySession {
   // go straight to the cleaning phase.
   private boolean trainingFinished = false;
 
-  private Set<String> registeredTasks = new HashSet<>();
+  private Set<String> registeredTasks = ConcurrentHashMap.newKeySet();
 
   private int numExpectedTasks = 0;
 
@@ -677,7 +676,7 @@ public void addRegisteredTask(String taskId) {
   }
 
   public void resetRegisteredTasks() {
-    registeredTasks = new HashSet<>();
+    registeredTasks = ConcurrentHashMap.newKeySet();
   }
 
   public int getNumRegisteredTasks() {