blackfish:

xs = xs.view(-1, 10 * 3 * 6) print(‘xs shape’, xs.shape) theta = self.fc_loc(xs)


this looks wrong. you need to keep the first dimension to be 1 (it’s the batch size), so dont give -1 in the first dimension.
See that at the end, theta shape has a batch size of 1728 which is wrong, because you started with batch size of 1