Moving data to GPU in collate_fn fails

I do recommend posting something on this thread. The devs seem to be taking quite an interest in it and they’ve been active there.

Shubh