diff --git a/.gitignore b/.gitignore
index 62bf21f..14dec7c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.idea
 .DS_Store
 /pkg/*
 
diff --git a/lib/generators/ratyrate/templates/average_cache_migration.rb b/lib/generators/ratyrate/templates/average_cache_migration.rb
index 27392ce..b69ce17 100644
--- a/lib/generators/ratyrate/templates/average_cache_migration.rb
+++ b/lib/generators/ratyrate/templates/average_cache_migration.rb
@@ -1,4 +1,4 @@
-class CreateAverageCaches < ActiveRecord::Migration
+class CreateAverageCaches < ActiveRecord::Migration[5.2]
 
   def self.up
     create_table :average_caches do |t|
@@ -13,6 +13,7 @@ def self.up
 
   def self.down
     drop_table :average_caches
+    remove_index :rating_caches, [:cacheable_id, :cacheable_type]
   end
 
 end
diff --git a/lib/generators/ratyrate/templates/migration.rb b/lib/generators/ratyrate/templates/migration.rb
index a15dc84..92efc2c 100644
--- a/lib/generators/ratyrate/templates/migration.rb
+++ b/lib/generators/ratyrate/templates/migration.rb
@@ -1,4 +1,4 @@
-class CreateRates < ActiveRecord::Migration
+class CreateRates < ActiveRecord::Migration[5.2]
 
   def self.up
       create_table :rates do |t|
@@ -15,6 +15,8 @@ def self.up
 
     def self.down
       drop_table :rates
+      remove_index :rates, :rater_id
+      remove_index :rates, [:rateable_id, :rateable_type]
     end
 
 end
\ No newline at end of file
diff --git a/lib/generators/ratyrate/templates/overall_average_migration.rb b/lib/generators/ratyrate/templates/overall_average_migration.rb
index 8aba69a..19152a1 100644
--- a/lib/generators/ratyrate/templates/overall_average_migration.rb
+++ b/lib/generators/ratyrate/templates/overall_average_migration.rb
@@ -1,4 +1,4 @@
-class CreateOverallAverages < ActiveRecord::Migration
+class CreateOverallAverages < ActiveRecord::Migration[5.2]
 
   def self.up
     create_table :overall_averages do |t|
diff --git a/lib/ratyrate/model.rb b/lib/ratyrate/model.rb
index ce8d1f0..b9ee178 100644
--- a/lib/ratyrate/model.rb
+++ b/lib/ratyrate/model.rb
@@ -64,26 +64,27 @@ def update_current_rate(stars, user, dimension)
   end
 
   def overall_avg(user)
-    # avg = OverallAverage.where(rateable_id: self.id)
-    # #FIXME: Fix the bug when the movie has no ratings
-    # unless avg.empty?
-    #   return avg.take.avg unless avg.take.avg == 0
-    # else # calculate average, and save it
-    #   dimensions_count = overall_score = 0
-    #   user.ratings_given.select('DISTINCT dimension').each do |d|
-    #     dimensions_count = dimensions_count + 1
-    #     unless average(d.dimension).nil?
-    #       overall_score = overall_score + average(d.dimension).avg
-    #     end
-    #   end
-    #   overall_avg = (overall_score / dimensions_count).to_f.round(1)
-    #   AverageCache.create! do |a|
-    #     a.rater_id = user.id
-    #     a.rateable_id = self.id
-    #     a.avg = overall_avg
-    #   end
-    #   overall_avg
-    # end
+    avg = OverallAverage.where(rateable_id: self.id)
+    #FIXME: Fix the bug when the movie has no ratings
+    unless avg.empty?
+      return avg.take.avg unless avg.take.avg == 0
+    else # calculate average, and save it
+      dimensions_count = overall_score = 0
+      user.ratings_given.select('DISTINCT dimension').each do |d|
+        dimensions_count = dimensions_count + 1
+        unless average(d.dimension).nil?
+          overall_score = overall_score + average(d.dimension).avg
+        end
+      end
+      overall_avg = (overall_score / dimensions_count).to_f.round(1)
+      AverageCache.create! do |a|
+        a.rater_id = user.id
+        a.rateable_id = self.id
+        a.rateable_type = self.class.to_s
+        a.avg = overall_avg
+      end
+      overall_avg
+    end
   end
 
   # calculates the movie overall average rating for all users