diff --git a/lib/generators/acts_as_follower_generator.rb b/lib/generators/acts_as_follower_generator.rb
index 75ce101..fcf5051 100644
--- a/lib/generators/acts_as_follower_generator.rb
+++ b/lib/generators/acts_as_follower_generator.rb
@@ -20,11 +20,23 @@ def self.next_migration_number(dirname)
   end
   
   def create_migration_file
-    migration_template 'migration.rb', 'db/migrate/acts_as_follower_migration.rb'
+    migration_template 'migration.rb', 'db/migrate/acts_as_follower_migration.rb', migration_version: migration_version
   end
   
   def create_model
     template "model.rb", File.join('app/models', "follow.rb")
   end
+
+  def rails5?
+    Rails.version.start_with? '5'
+  end
+  def migration_version
+    if rails5?
+      "[#{Rails::VERSION::MAJOR}.#{Rails::VERSION::MINOR}]"
+    else
+      "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      
+    end
+  end
   
 end
diff --git a/lib/generators/templates/migration.rb b/lib/generators/templates/migration.rb
index 0b39c01..1043117 100644
--- a/lib/generators/templates/migration.rb
+++ b/lib/generators/templates/migration.rb
@@ -1,4 +1,4 @@
-class ActsAsFollowerMigration < ActiveRecord::Migration
+class ActsAsFollowerMigration < ActiveRecord::Migration<%= migration_version %>
   def self.up
     create_table :follows, force: true do |t|
       t.references :followable, polymorphic: true, null: false