If self.model is wrapped (DataParallel, torch.compile, or future PolicyValueNetwork path), forward() can return a tuple (logits, value) or even a dict. Add an explicit unwrap immediately after the forward call in both predict() and predict_batch() so every downstream temperature/softmax operation always receives a plain tensor.