diff --git a/data_utils.py b/data_utils.py
index 86094ca..c40fcb7 100644
--- a/data_utils.py
+++ b/data_utils.py
@@ -1,120 +1,120 @@
 import numpy as np
 import torch
 import random
 
 def computeVelocity(coords, timestep=4):
     '''
     Compute speed of pedestrian at each timestep given its coordinates during the sequence
     timestep = 4 seconds by default (observations each 10 frames with a 2.5 framerate)
     Returns velocity vector (vx, vy)
     '''
     npos = coords.size()[0] #number of observed positions
     velocity = np.zeros((npos,2))
     x = coords[:,0]
     y = coords[:,1]
 
     for i in range(1, npos):
         velocity[i,0] =  (x[i]-x[i-1])/timestep
         velocity[i,1] =  (y[i]-y[i-1])/timestep
 
     velocity = torch.from_numpy(velocity)
     return velocity
 
 def getTargets(coords):
     '''
     returns a target array for given coordinates for supervised training
     '''
     targets = coords[1:]
 
     return targets
 
 def flipPosition(traindata,method):
     '''
     Flips given coordinates of all pedestrian in traindata
     method:
     1: change signs of x position (vertical symmetry)
     2: change signs of y position (horizontal symmetry)
     3: change signs of x and y positions
     '''
     traindata_new=traindata
     for p in range(0, len(traindata)): #loop of over each pedestrian
             #print(traindata[p].coords)
         if traindata[p].neighbors is None:
             continue
         #print('coords unflipped')
         #print(traindata[p].coords)
         if method == 1:
             #change signs of x position
             #print('befor')
             traindata[p].coords[:,0] = - traindata[p].coords[:,0]
             traindata[p].neighbors[:,2] = - traindata[p].neighbors[:,2]
             #print(traindata[p].coords)
         elif method == 2:
             #change signs of y position
             traindata[p].coords[:,1] = - traindata[p].coords[:,1]
             traindata[p].neighbors[:,3] = - traindata[p].neighbors[:,3]
 
             #coords[:,1] = -coords[:,1]
         else:
             #change both signs
             traindata[p].coords = - traindata[p].coords
             traindata[p].neighbors[:,2:4] = - traindata[p].neighbors[:,2:4]
 
     return traindata
 
 def getCoords(velocity, init_coords, timestep = 4):
     '''
     Compute x and y coordinates given a velocity array and the initial coordinates
     '''
     count = 0;
     x_old = init_coords[:,0]
     y_old = init_coords[:,1]
     x_pos = []
     y_pos = []
 
     for i in range(velocity.size()[0]):
 
         x_new = velocity[i,0]*timestep + x_old
         y_new = velocity[i,1]*timestep + y_old
 
         x_pos += [x_new]
         y_pos += [y_new]
 
         x_old = x_new
         y_old = y_new
 
 
     x_pos = torch.cat(x_pos,0)
     y_pos = torch.cat(y_pos,0)
 
     return x_pos, y_pos
 
 def getOccupancyTensor(pedestrian):
     '''
     Compute the occupancy map pooling (Alahi et al.)
     pedestrian : pedestrian object
     '''
 
     current_frames = pedestrian.frames #frames in which the current pedestrian appears
     ped_coords = pedestrian.coords
     neighbors = pedestrian.neighbors
 
     if neighbors is None:
 
         occupancy = None
 
     else:
 
         neigh_frames = neighbors[:,0] #frames in which the neighbors appear
         neigh_coords = neighbors[:,2:4] #coordinates of the neighbors
 
         occupancy = []
         for i, frame in enumerate(current_frames):
-            pooled_x = torch.sum(neigh_coords[(neigh_frames == frame).nonzero(),0]-ped_coords[i, 0])
-            pooled_y = torch.sum(neigh_coords[(neigh_frames == frame).nonzero(),1]-ped_coords[i, 1])
+            pooled_x = torch.sum(neighbors[(neighbors[:,0] == frame).nonzero(),2]-ped_coords[i, 0])
+            pooled_y = torch.sum(neighbors[(neighbors[:,0] == frame).nonzero(),3]-ped_coords[i, 1])
 
             occupancy += [(pooled_x, pooled_y)]
 
         occupancy = torch.cat(occupancy, 0)
 
     return occupancy