BUG: Scheduling while Atomic - Kernel Driver on TX2

I have a kernel driver, where there is a section, in which a spin_lock() is applied and a wait_event_interruptible_timeout() within this spin_lock.
When I insmod the driver for this card, Scheduling while atomic BUG is coming and the Jetson TX2 is crashing (restarting) abruptly. The same driver works fine on an x86.
So is this anything related to the architecture?
Can anyone shed some light on this issue?

spin_lock(...);
....
....
    wait_event_interruptible_timeout(....);
....
....
spin_unlock(....);

Did you try on x86 with K3.18 too?

Actually, I am using Jetpack 3.3 with K4.4. Tried that kernel on x86.