I made some modifications to the LSTMCell function in nn_functions\rnn.py which works as intended when I run the network on my CPU. But when I use cuda the network seems to skip the LSTMCell function.
The network will still run when I use cuda, but when I throw in a print(‘working’) into the LSTMCell function, it doesn’t print, implying that the LSTMCell function gets skipped when using cuda.
In which other files must I make modifications for my changes to carry over when the network is run on GPU?
Illustration:
def LSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
    if input.is_cuda: 
        igates = F.linear(input, w_ih)
        hgates = F.linear(hidden[0], w_hh)
        state = fusedBackend.LSTMFused.apply
        return state(igates, hgates, hidden[1]) if b_ih is None else state(igates, hgates, hidden[1], b_ih, b_hh)
        #I can comment the above section out entirely and the the network still runs

    hx, cx = hidden
    gates = F.linear(input, w_ih, b_ih) + F.linear(hx, w_hh, b_hh) 
    ingate, forgetgate, cellgate, outgate, choosegate = gates.chunk(5, 1)  

    choosegate_1 = (F.sigmoid(choosegate)).round()
    choosegate_2 = (1-choosegate_1)
    cellgate_1 = F.tanh(cellgate)      
    cellgate_2 = F.relu(cellgate)   

    ingate = F.sigmoid(ingate)
    forgetgate = F.sigmoid(forgetgate)
    outgate = F.sigmoid(outgate)
    cellgate = ((choosegate_1 * cellgate_1) + (choosegate_2 * cellgate_2))

    cy = (forgetgate * cx) + ingate*cellgate
    hy = outgate * F.tanh(cy)
    return hy, cy

My modifications of the LSTMCell work properly when run without cuda, but it seems that cuda may have its own construction of the LSTMCell that I cannot find.
If someone could point out where cuda runs the LSTMCell layers through activation functions, or more generally constructs the entire LSTMCell, I would be greatly appreciative!
Thanks!!