Bug fix

Previous commit had some problems. Fixed now.