Fix UT
Merge pull request #281 from asadoughi/master

Fix UT