I was using shapenet/pvcnn.py with the following modified forward():
    def forward(self, inputs):
        # inputs : [B, in_channels + S, N]
        features = inputs[:, :self.in_channels, :] 
        num_points = features.size(-1)

        # prepare feature extraction for NST
        feats = {}

        coords = features[:, :3, :]
        out_features_list = [one_hot_vectors]
        for i in range(len(self.point_features)):
            features, _ = self.point_features[i]((features, coords)) # takes 6d (feat), 3d (coord) array
            feats["point_feature_{}".format(i)] = features

        # max pool out features to global feature
        global_feature = features.max(dim=-1, keepdim=True).values.repeat([1, 1, num_points])
        feats["global_feature_{}".format(i)] = global_feature

        return feats


The input is a [1, 6, N] tensor, where N is the number of 3d points