Problem with BERT, XLNET Transformer Pipline and GPU acceleration with Flask API

I am beginner in the field of machine learning and would really appreciate any help. Thank you in advance.

I have learned XLNet and BERT Transformer models, which both solve the text classification problem. Both artificial networks work well and I don’t see any problem while using them on my GPU and CPU as well directly from my python code.

Few weeks ago I decided to implement a simple API to access the learned models and I have implemented a simple Flask application that uses Worker and runs in both synchronous and asynchronous modes (https://realpython.com/flask-by-example-implementing-a-redis-task-queue/). The application works good and I am able to access the models to let them classify the input text (the text is send using POST request in JSON format), but there is one problem which I am not able to solve.

The API works without any problem only during the synchronous mode (both GPU and CPU versions). The problem appears when I tried to put the same code from my CPU to my GPU and run it again asynchronously. The problem appears when I try to access model’s predictions (its result, see ret variable).

This happens for both the Transformer Pipline and for my bert model in eval() mode. In case when I try to access model's result the Python returned me (see the exec_job function) ret is None; TypeError: 'NoneType' object is not subscriptable and I am not able to access ret['status'], because the ret is None.

When I delete the parts of code putting my model and created tensors on GPU (I just deleted .to(_configuration.GPU_DEVICE_ID) parts) - everything works fine (but in this case my model is located on CPU and the evaluation takes much more time).

# This code runs in __init__
self._model.to(_configuration.GPU_DEVICE_ID) 

# These lines are executed during evaluation         
inputs = torch.tensor(input_ids).to(_configuration.GPU_DEVICE_ID)
segments = torch.tensor(segments).to(_configuration.GPU_DEVICE_ID)
masks = torch.tensor(attention_masks).to(_configuration.GPU_DEVICE_ID)

I tied to debug my code and notice that the function self._process_tokenized_text(sentence) executed successfully but the code after GPU parts not executed at all and even my print() code does not do anything.

Simplified code below.

Asynchronous function which runs the function exec_job

    @app.route('/api/text/check-json-async', methods=['POST'])
    @auth.login_required
    def check_json_async():
        request_data = request.get_json()
        text=request_data['data']
        job = q.enqueue_call(
            func='aijob.exec_job', args=(text,True,), result_ttl=60
        )
        global callbackURL
        ret={"text":text,"callbackURL":callbackURL,"jobID":job.get_id(),"status":"queued"}
        return jsonify(ret)

exec_job function which calls _aiprocess function and access the AI core

    # The function whitch publish the main functionality of this layer
    def exec_job(text, async_parm):
        # Evaluate text
        ret = _aiprocess(text)
        # Save classificated result --> THIS CODE FAILS ONLY ON GPU MODE BECAUSE the RET is NONE
        _savePredictedResult(ret['status'] == STATUS, ret)
        if async_parm is True:
            import pickle
            f = open('callbackURL.pckl', 'rb')
            callbackURL = pickle.load(f)
            f.close()
            print(ret)
            print(callbackURL)
            requests.post(callbackURL, json=ret)
        return ret

The AI core the _aiprocess function

   def _aiprocess(self, sentence):
        if sentence is not None:
            #function working on CPU successfully returns the result
            input_ids, attention_masks, segments = self._process_tokenized_text(sentence)
            # print works as well
            print(input_ids, attention_masks, segments)

            inputs = torch.tensor(input_ids).to(_configuration.GPU_DEVICE_ID)
            segments = torch.tensor(segments).to(_configuration.GPU_DEVICE_ID)
            masks = torch.tensor(attention_masks).to(_configuration.GPU_DEVICE_ID)
            # This code do nothing and when I try to access the predictions array there is nothing but only in GPU mode
            print(inputs)
            # Tracking variables 
            predictions = []
            
            # Predict 
            with torch.no_grad():
                outputs = self._model(input_ids = inputs,token_type_ids=segments, input_mask = masks)
                for tmptuple in outputs[:1]:
                    tmp_array = torch.nn.functional.softmax(tmptuple, dim=1).cpu().numpy()
                    for result in tmp_array:
                        predictions.append(self._resultPrediction(result))                                   
            return predictions[0]

Do you have any idea why this is happening? Thank you so much and I am sorry for my bad English.