I will take a closer look later but the issue is that you need to use Theano’s functions for Boolean comparison instead of just vanilla python if/else statements. Check out theano.tensor.switch. You will also want to use theano.scan for your for loop, but maybe you can get by without it.
1 Like