Added forgotten code this needs to be verified

I havent been able to test this due to time constraints, simple hack to
trawl back two lines, although I may not have compled it fully.