You can set the batch size manually using trainer.prediction_loop()
Instead of using trainer.predict(test_dataset), you can use torch DataLoader for trainer.prediction_loop(). Thus, you might change
from
raw_pred, _, _ = trainer.predict(test_dataset)

into:
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
raw_pred, _, _ = trainer.prediction_loop(test_loader, description="prediction")