Parallel patch embed for faster execution

[rejuvyesh]

We can use the trick of grouped convolutions to perform patch embedding of different variables in parallel in a single GPU kernel call instead of serially in a for loop.

TODO:

• Correctly converting the non-parallel patch embed checkpoints

[rejuvyesh]

Some rough timing:

parallel_patch_embed: True, img_size: [32,64]
forward time: 0.009
---
parallel_patch_embed: False, img_size: [32,64]
forward time: 0.021
---
parallel_patch_embed: True, img_size: [128, 256]
forward time: 0.097
---
parallel_patch_embed: False, img_size: [128, 256]
forward time: 0.109

parallel_patch_embed: True, img_size: [32, 64]
backward time: 0.034
---
parallel_patch_embed: False, img_size: [32, 64]
backward time: 0.052