with_cpu
Add tests for Linear layer and update matmul usage
Wheel build configuration
Fix training issues