fix a potential issue

wbo4958 · Jan 8, 2025 · ea6d4d8 · ea6d4d8
1 parent 20a36f8
commit ea6d4d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/...ges/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/spark/ExternalMemory.scala b/...ges/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/spark/ExternalMemory.scala
@@ -206,7 +206,7 @@ private[scala] class ExternalMemoryIterator(val input: Iterator[Table],
   override def next(): ColumnBatch = {
     inputNextIsCalled = true
     withResource(new GpuColumnBatch(iter.next())) { batch =>
-      if (iter == input) {
+      if (iter.eq(input)) {
         externalMemory.cacheTable(batch.table)
       }
       new CudfColumnBatch(