Android 5.0 Lollipop notification full sound doesn't play

I have an alarm app that plays a sound (alarm-like continuous audio) when the alarm goes off. Unfortunately, in Lollipop the sound doesn't play in full, instead it gets stopped after a couple of seconds. But if the phone is connected to power, this doesn't happen and the sound actually is played fully. The code works fine on previous versions of Android. Can anyone please help? Here's my code for the notification:

NotificationCompat.Builder mBuilder = new NotificationCompat.Builder(context)
                .setAutoCancel(true)
                .setPriority(Integer.MAX_VALUE)
                .setContentTitle(someTitle)
                .setWhen(now)
                .setIcon(R.drawable.some_icon);

Notification notif = mBuilder.build();

if(Build.VERSION.SDK_INT >= 21) {
    notif.sound = audioFileUri;
    notif.category = Notification.CATEGORY_ALARM;

    AudioAttributes.Builder attrs = new AudioAttributes.Builder();
    attrs.setContentType(AudioAttributes.CONTENT_TYPE_SONIFICATION);
    attrs.setUsage(useAlarm ? AudioAttributes.USAGE_ALARM : AudioAttributes.USAGE_NOTIFICATION_EVENT);
    notif.audioAttributes = attrs.build();
} else  {
    mBuilder.setSound(audioFileUri, useAlarm ? AudioManager.STREAM_ALARM : AudioManager.STREAM_NOTIFICATION);
    notif = mBuilder.build();
}

NotificationManager mNotificationManager = (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
mNotificationManager.notify(NOTIFICATION_ID, notif);