Added fix to run predict just on head node
@@ -300,6 +300,8 @@ class BERTSequenceClassifier:
1darray, dict(1darray, 1darray, ndarray): Predicted classes and target labels or
a dictionary with classes, target labels, probabilities) if probabilities is True.
"""
+
+ if hvd.rank() == 0:
device = get_device("cpu" if num_gpus == 0 else "gpu")
self.model = move_to_device(self.model, device, num_gpus)
