Considerations of model definitions when moving from Tensorflow to PyTorch
I've just recently switched to PyTorch after getting frustrated in debugging tf and understand that it is equivalent to coding in numpy almost completely. My question is what are the permitted python aspects we can use in a PyTorch model (to be put completely on GPU) eg. if-else has to be implemented as follows in tensorflow

a = tf.Variable([1,2,3,4,5], dtype=tf.float32)
b = tf.Variable([6,7,8,9,10], dtype=tf.float32)
p = tf.placeholder(dtype=tf.float32)
ps = tf.placeholder(dtype=tf.bool)

li = [None]*5
li_switch = [True, False, False, True, True]

for i in range(5):
    li[i] = tf.Variable(tf.random.normal([5]))

sess = tf.Session()
sess.run(tf.global_variables_initializer())

def func_0():
    return tf.add(a, p)
def func_1():
    return tf.subtract(b, p)

with tf.device('GPU:0'):
    my_op = tf.cond(ps, func_1, func_0)

for i in range(5):
    print(sess.run(my_op, feed_dict={p:li[i], ps:li_switch[i]}))


How would the structure change in pytorch for the above code? How to place the variables and ops above on GPU and parallelize the list inputs to our graph in pytorch?

In pytorch, the code can be written like the way normal python code is written.

CPU

import torch
a = torch.FloatTensor([1,2,3,4,5])
b = torch.FloatTensor([6,7,8,9,10])
cond = torch.randn(5)

for ci in cond:
    if ci &gt; 0:
        print(torch.add(a, 1))
    else:
        print(torch.sub(b, 1))


GPU

Move the tensors to GPU like this:

a = torch.FloatTensor([1,2,3,4,5]).to('cuda')
b = torch.FloatTensor([6,7,8,9,10]).to('cuda')
cond = torch.randn(5).to('cuda')

import torch.nn as nn

class Cond(nn.Module):
    def __init__(self):
        super(Cond, self).__init__()

    def forward(self, cond, a, b):
        result =  torch.empty(cond.shape[0], a.shape[0]).cuda()
        for i, ci in enumerate(cond):
            if ci &gt; 0:
                result[i] = torch.add(a, 1)
            else:
                result[i] = torch.sub(b, 1)

        return result

cond_model = Cond().to('cuda')
output = cond_model(cond, a, b)


https://pytorch.org/tutorials/beginner/blitz/tensor_tutorial.html#cuda-tensors
