Feature/#21

Fix #21 problem