Fixed optim update error with non-contiguous grads/params

[Edenzzzz]

Fixes #1185 Non-contiguous params/gradients resulting from torch.chunk and all_gather etc. are ubiquitous in distributed training frameworks such as ZeRO. This avoids update errors as the C++ kernels assume row-major inputs.

cc @Titus-von-Koeller