Fix type check error
Fix type check error
Fix type check error
Fix type check errors and minor cleanup