trying to fix the bug of dimension mismatch