When things are taking a long time it’s usually because there’s some unexpected broadcasting going on.
I suggest to comment out everything related to sampling, and look at p.eval() and obs.eval().
The .eval() method runs the computational graph associated with a variable and lets you see a sample output. My guess is that you will find that one of these is not (n,1), but (n,n) or something.