Fixed the issue #521

Hi, this PR fixed the issue #521. Please take a look. Thanks!