Fix #231

Comment From: tzolov

@Omkar-Shetkar thanks for the contribution.

Comment From: tzolov

Rebased, cleaned and squashed at 0b87df07e5a4da49df225e51b9dc08c0fd795968