The UnsafeViewBackward shouldn’t be problematic as explained here.
However, note that the linear layer will behave differently for a 2D and 3D input.
While a 2D input should be the “standard” approach, a 3D input would be treated as a sequence (in dim1) or inputs.