Fix the strides warning

[ShiromiyaG]

This PR is to try to fix the stride warning at the start of the training. Problems: the current solution slows down training and give cuda out of memory. To make the training faster and solve the cuda out of memory problem, just set create_graph to false, but it's still slower than the original code