diff --git a/sentiment-analysis/sentiment-analysis.py b/sentiment-analysis/sentiment-analysis.py
index 3fa0094..3889cf8 100644
--- a/sentiment-analysis/sentiment-analysis.py
+++ b/sentiment-analysis/sentiment-analysis.py
@@ -1,30 +1,56 @@
 import os
 import pandas as pd
-import torch
-from transformers import DistilBertTokenizerFast
-from torch.utils.data import DataLoader
-from transformers import DistilBertForSequenceClassification
 from tqdm import tqdm
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import (
+    DistilBertTokenizerFast,
+    DistilBertForSequenceClassification
+)
 
 # Importing the csv dataset.
 case = '...'
-df_data = pd.read_csv(f'../datasets/{case}.csv')
+df_data = pd.read_csv(f'{case}.csv')
 print(len(df_data))
 
-# Preparing the dataset for the network this includes tokenization, encoding and creating dataloaders.
+# Set device cuda
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print('Using:', device)
+
+# Define model
 model_name = 'distilbert-base-uncased'
 tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
 
-encodings = tokenizer(df_data['text'].tolist(), truncation=True, padding=True, return_tensors='pt')
+# Tokenize and encode the text data
+encodings = tokenizer(
+    df_data['text'].tolist(), 
+    truncation=True, 
+    padding=True, 
+    return_tensors='pt'
+)
 
-dataset = torch.utils.data.TensorDataset(
+# Create tensor dataset
+dataset = TensorDataset(
     encodings['input_ids'], 
     encodings['attention_mask']
 )
-dataloader = DataLoader(dataset, batch_size=45, shuffle=False)
 
-# Loading the model.
-model= torch.load('models/bert-aclimdb.pth')
+# Create data loader
+if device == 'cuda':
+    en_pin_memory = True
+else:
+    en_pin_memory = False
+dataloader = DataLoader(
+    dataset, 
+    batch_size=53, 
+    shuffle=False,
+    pin_memory=en_pin_memory
+)
+
+# Loading the model
+model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)
+model.load_state_dict(torch.load('...', map_location=device))
+model.to(device)
 
 # Using the model to perform sentiment analysis on the dataset.
 model.eval()
@@ -32,7 +58,7 @@ list_predicted_scores = []
 
 for batch in tqdm(dataloader):
     with torch.no_grad():
-        input_ids, attention_mask = batch
+        input_ids, attention_mask = [x.to(device) for x in batch]
 
         # Obtaining the sentiment score.
         output = model(input_ids=input_ids, attention_mask=attention_mask)
@@ -45,4 +71,4 @@ for batch in tqdm(dataloader):
 df_data['s'] = list_predicted_scores
 
 # Exporting the dataframe to a csv dataset.
-df_data.to_csv(f'../datasets/{case}-s.csv', index=False)
\ No newline at end of file
+df_data.to_csv(f'{case}-s.csv', index=False)